From 8b326137c6f4f8f5bc87fe67c895cb17564b1f6e Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Thu, 18 Aug 2022 16:13:18 +0800 Subject: [PATCH] Add a new attribute status to ExecutorInfo file 1. Add a new attribute status to ExecutorInfo file 2. Update running status to the executor infofile. --- dvc/repo/experiments/__init__.py | 4 +++- dvc/repo/experiments/executor/base.py | 24 ++++++++++++++++++++++-- dvc/repo/experiments/executor/local.py | 4 +++- dvc/repo/experiments/queue/base.py | 7 +++++++ dvc/repo/experiments/queue/workspace.py | 5 ++++- 5 files changed, 39 insertions(+), 5 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index b0cec7380..fb9ae99ff 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -16,7 +16,7 @@ from .exceptions import ( InvalidExpRefError, MultipleBranchError, ) -from .executor.base import BaseExecutor, ExecutorInfo +from .executor.base import BaseExecutor, ExecutorInfo, TaskStatus from .queue.base import BaseStashQueue, QueueEntry from .queue.celery import LocalCeleryQueue from .queue.tempdir import TempDirQueue @@ -463,6 +463,8 @@ class Experiments: if rev == "workspace": # If we are appending to a checkpoint branch in a workspace # run, show the latest checkpoint as running. + if info.status > TaskStatus.RUNNING: + return result last_rev = self.scm.get_ref(EXEC_BRANCH) if last_rev: result[last_rev] = info.asdict() diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 4595507a9..7c53d2749 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -4,6 +4,7 @@ import pickle from abc import ABC, abstractmethod from contextlib import contextmanager from dataclasses import asdict, dataclass +from enum import IntEnum from functools import partial from typing import ( TYPE_CHECKING, @@ -49,10 +50,9 @@ if TYPE_CHECKING: from scmrepo.git import Git from dvc.repo import Repo + from dvc.repo.experiments.stash import ExpStashEntry from dvc.stage import PipelineStage - from ..base import ExpStashEntry - logger = logging.getLogger(__name__) @@ -66,6 +66,15 @@ class ExecutorResult(NamedTuple): force: bool +class TaskStatus(IntEnum): + PENDING = 0 + PREPARING = 1 + RUNNING = 2 + SUCCESS = 3 + FAILED = 4 + REVOKED = 5 + + @dataclass class ExecutorInfo: git_url: str @@ -79,6 +88,7 @@ class ExecutorInfo: result_hash: Optional[str] = None result_ref: Optional[str] = None result_force: bool = False + status: TaskStatus = TaskStatus.PENDING @classmethod def from_dict(cls, d): @@ -138,6 +148,7 @@ class BaseExecutor(ABC): root_dir: str, dvc_dir: str, baseline_rev: str, + status: TaskStatus, wdir: Optional[str] = None, name: Optional[str] = None, location: Optional[str] = None, @@ -151,6 +162,7 @@ class BaseExecutor(ABC): self.baseline_rev = baseline_rev self.location: str = location or self.DEFAULT_LOCATION self.result = result + self.status = status @abstractmethod def init_git( @@ -199,6 +211,7 @@ class BaseExecutor(ABC): dvc_dir=self.dvc_dir, name=self.name, wdir=self.wdir, + status=self.status, **result_dict, ) @@ -220,6 +233,7 @@ class BaseExecutor(ABC): root_dir=info.root_dir, dvc_dir=info.dvc_dir, baseline_rev=info.baseline_rev, + status=info.status, name=info.name, wdir=info.wdir, result=result, @@ -247,6 +261,7 @@ class BaseExecutor(ABC): root_dir=root_dir, dvc_dir=relpath(repo.dvc_dir, repo.scm.root_dir), baseline_rev=entry.baseline_rev, + status=TaskStatus.PREPARING, name=entry.name, wdir=relpath(os.getcwd(), repo.scm.root_dir), **kwargs, @@ -416,6 +431,7 @@ class BaseExecutor(ABC): exp_ref: Optional["ExpRefInfo"] = None repro_force: bool = False + info.status = TaskStatus.PREPARING if infofile is not None: info.dump_json(infofile) @@ -472,6 +488,9 @@ class BaseExecutor(ABC): info.name, repro_force or checkpoint_reset, ) + info.status = TaskStatus.RUNNING + if infofile is not None: + info.dump_json(infofile) stages = dvc_reproduce( dvc, *args, @@ -495,6 +514,7 @@ class BaseExecutor(ABC): info.result_hash = exp_hash info.result_ref = ref info.result_force = repro_force + info.status = TaskStatus.SUCCESS if infofile is not None: info.dump_json(infofile) diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index e977cd8f6..5d64345e0 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -19,7 +19,7 @@ from ..refs import ( EXEC_MERGE, EXEC_NAMESPACE, ) -from .base import EXEC_TMP_DIR, BaseExecutor +from .base import EXEC_TMP_DIR, BaseExecutor, TaskStatus if TYPE_CHECKING: from scmrepo.git import Git @@ -80,6 +80,8 @@ class TempDirExecutor(BaseLocalExecutor): DulwichRepo.init(os.fspath(self.root_dir)) + self.status = TaskStatus.PREPARING + with self.set_exec_refs(scm, stash_rev, entry): refspec = f"{EXEC_NAMESPACE}/" push_refspec(scm, self.git_url, refspec, refspec) diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index d228e5929..6bfadd440 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -580,9 +580,16 @@ class BaseStashQueue(ABC): exp.repo, stash_entry, **kwargs ) + infofile = exp.celery_queue.get_infofile_path(stash_rev) + executor.info.dump_json(infofile) + executor.init_git( exp.repo.scm, stash_rev, stash_entry, branch=stash_entry.branch ) + + infofile = exp.celery_queue.get_infofile_path(stash_rev) + executor.info.dump_json(infofile) + executor.init_cache(exp.repo, stash_rev) return executor diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index ff92dfcb1..82f80d241 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -9,7 +9,7 @@ from dvc.exceptions import DvcException from dvc.utils.fs import remove from ..exceptions import ExpQueueEmptyError -from ..executor.base import BaseExecutor, ExecutorResult +from ..executor.base import BaseExecutor, ExecutorResult, TaskStatus from ..executor.local import WorkspaceExecutor from ..refs import EXEC_BRANCH from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult @@ -110,10 +110,13 @@ class WorkspaceQueue(BaseStashQueue): ) except CheckpointKilledError: # Checkpoint errors have already been logged + executor.status = TaskStatus.FAILED return {} except DvcException: + executor.status = TaskStatus.FAILED raise except Exception as exc: + executor.status = TaskStatus.FAILED raise DvcException( f"Failed to reproduce experiment '{rev[:7]}'" ) from exc