From 44820e74a3e4cfb7201d37b00508571feb0d96df Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Wed, 27 Apr 2022 15:19:23 +0900 Subject: [PATCH] exp: wait for task collection when getting exp result --- dvc/repo/experiments/__init__.py | 22 ++++++++----------- dvc/repo/experiments/executor/base.py | 7 ++++++ dvc/repo/experiments/queue/base.py | 5 ++++- dvc/repo/experiments/queue/local.py | 17 +++++++++------ dvc/repo/experiments/queue/tasks.py | 31 +++++++++++++-------------- 5 files changed, 45 insertions(+), 37 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 404054dfa..5d507fe71 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -177,22 +177,18 @@ class Experiments: 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) + result = self.celery_queue.get_result(entry) + if result is None or result.exp_hash is None: + name = entry.name or entry.stash_rev[:7] + failed.append(name) + elif result.ref_info: + exp_rev = self.scm.get_ref(str(result.ref_info)) + results[exp_rev] = result.exp_hash 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) + if results: + self._log_reproduced((rev for rev in results), True) return results def _log_reproduced(self, revs: Iterable[str], tmp_dir: bool = False): diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 9549d1b75..64d5fde12 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -73,6 +73,7 @@ class ExecutorInfo: dvc_dir: str name: Optional[str] = None wdir: Optional[str] = None + collected: bool = False result_hash: Optional[str] = None result_ref: Optional[str] = None result_force: bool = False @@ -102,6 +103,12 @@ class ExecutorInfo: with modify_json(filename) as d: d.update(self.asdict()) + @classmethod + def load_json(cls, filename: str) -> "ExecutorInfo": + from dvc.utils.serialize import load_json + + return cls.from_dict(load_json(filename)) + _T = TypeVar("_T", bound="BaseExecutor") diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index c48a89307..620da7808 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -165,7 +165,10 @@ class BaseStashQueue(ABC): @abstractmethod def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]: - """Return result of the specified item.""" + """Return result of the specified item. + + This method blocks until the specified item has been collected. + """ def _stash_exp( self, diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index a79b48098..05c9e5770 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -1,7 +1,7 @@ import hashlib -import json import logging import os +import time from collections import defaultdict from typing import ( TYPE_CHECKING, @@ -155,12 +155,15 @@ class LocalCeleryQueue(BaseStashQueue): 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 + while True: + try: + executor_info = ExecutorInfo.load_json(infofile) + if executor_info.collected: + return executor_info.result + except FileNotFoundError: + # Infofile will not be created until execution begins + pass + time.sleep(1) class WorkspaceQueue(BaseStashQueue): diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index 7ba59101b..9104e39f2 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -1,11 +1,9 @@ -import json -import os from typing import Any, Dict from celery import chain, shared_task from celery.utils.log import get_task_logger -from dvc.utils.fs import makedirs, remove +from dvc.utils.fs import remove from ..executor.base import ExecutorInfo from ..executor.local import TempDirExecutor @@ -28,9 +26,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None: ) proc = repo.experiments.celery_queue.proc infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev) - makedirs(os.path.dirname(infofile), exist_ok=True) - with open(infofile, "w", encoding="utf-8") as fobj: - json.dump(executor.info.asdict(), fobj) + executor.info.dump_json(infofile) cmd = ["dvc", "exp", "exec-run", "--infofile", infofile] # schedule execution + cleanup @@ -57,19 +53,22 @@ def collect_exp( entry = QueueEntry.from_dict(entry_dict) repo = Repo(entry.dvc_root) - with open(infofile, encoding="utf-8") as fobj: - executor_info = ExecutorInfo.from_dict(json.load(fobj)) + executor_info = ExecutorInfo.load_json(infofile) logger.debug("Collecting experiment info '%s'", str(executor_info)) executor = TempDirExecutor.from_info(executor_info) exec_result = executor_info.result - if exec_result is not None: - results = BaseStashQueue.collect_executor( - repo.experiments, executor, exec_result - ) - for rev in results: - logger.debug("Collected experiment '%s'", rev[:7]) - else: - logger.debug("Exec result was None") + try: + if exec_result is not None: + results = BaseStashQueue.collect_executor( + repo.experiments, executor, exec_result + ) + for rev in results: + logger.debug("Collected experiment '%s'", rev[:7]) + else: + logger.debug("Exec result was None") + finally: + executor_info.collected = True + executor_info.dump_json(infofile) @shared_task(ignore_result=True)