diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 987ec9d0d..f46a09a78 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -144,13 +144,7 @@ class Experiments: ) return [entry.stash_rev] if tmp_dir: - self.celery_queue.spawn_worker() - # wait for task execution to start - while not self.celery_queue.proc.get(entry.stash_rev): - time.sleep(1) - for line in self.celery_queue.proc.follow(entry.stash_rev): - ui.write(line, end="") - return + return self.reproduce_celery(entries=[entry]) if machine: # TODO: decide how to handle queued remote execution raise NotImplementedError @@ -167,14 +161,41 @@ class Experiments: return last_applied return None - def reproduce_queued(self, **kwargs): + def reproduce_celery( + self, entries: Optional[Iterable[QueueEntry]] = None, **kwargs + ) -> Dict[str, str]: + results: Dict[str, str] = {} + if entries is None: + entries = list(self.celery_queue.iter_queued()) + if not entries: + return results + self.celery_queue.spawn_worker() - # wait for task execution to start - while not first(self.celery_queue.iter_active()): - time.sleep(1) - for entry in self.celery_queue.iter_active(): + failed = [] + for entry in entries: + # wait for task execution to start + while not self.celery_queue.proc.get(entry.stash_rev): + time.sleep(1) for line in self.celery_queue.proc.follow(entry.stash_rev): ui.write(line, end="") + # wait for task collection to complete + while True: + result = self.celery_queue.get_result(entry) + if result is not None: + if result.exp_hash is None: + name = entry.name or entry.stash_rev[:7] + failed.append(name) + break + exp_rev = self.scm.get_ref(str(result.ref_info)) + if exp_rev is not None: + results[exp_rev] = result.exp_hash + break + time.sleep(1) + if failed: + names = ", ".join(name for name in failed) + ui.error(f"Failed to reproduce experiment(s) '{names}'") + self._log_reproduced((rev for rev in results), True) + return results def _log_reproduced(self, revs: Iterable[str], tmp_dir: bool = False): names = [] diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 94c878b3a..b3e0f9e72 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -162,6 +162,10 @@ class BaseStashQueue(ABC): def reproduce(self) -> Mapping[str, Mapping[str, str]]: """Reproduce queued experiments sequentially.""" + @abstractmethod + def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]: + """Return result of the specified item.""" + def _stash_exp( self, *args, diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index 1a88ce655..dc20f098a 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -1,7 +1,16 @@ +import json import logging import os from collections import defaultdict -from typing import TYPE_CHECKING, Dict, Generator, List, Mapping, NamedTuple +from typing import ( + TYPE_CHECKING, + Dict, + Generator, + List, + Mapping, + NamedTuple, + Optional, +) from funcy import cached_property, first from kombu.message import Message @@ -13,6 +22,7 @@ from ..executor.base import ( EXEC_PID_DIR, EXEC_TMP_DIR, BaseExecutor, + ExecutorInfo, ExecutorResult, ) from ..executor.local import WorkspaceExecutor @@ -140,6 +150,15 @@ class LocalCeleryQueue(BaseStashQueue): def reproduce(self) -> Mapping[str, Mapping[str, str]]: raise NotImplementedError + def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]: + infofile = self.get_infofile_path(entry.stash_rev) + try: + with open(infofile, encoding="utf-8") as fobj: + executor_info = ExecutorInfo.from_dict(json.load(fobj)) + except FileNotFoundError: + return None + return executor_info.result + class WorkspaceQueue(BaseStashQueue): def put(self, *args, **kwargs) -> QueueEntry: @@ -253,3 +272,6 @@ class WorkspaceQueue(BaseStashQueue): logger.debug("Collected experiment '%s'.", exp_rev[:7]) results[exp_rev] = exec_result.exp_hash return results + + def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]: + raise NotImplementedError