diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 6047f27a5..8a5e4e1cf 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -116,6 +116,10 @@ class BaseStashQueue(ABC): def iter_queued(self) -> Generator[QueueEntry, None, None]: """Iterate over items in the queue.""" + @abstractmethod + def iter_active(self) -> Generator[QueueEntry, None, None]: + """Iterate over items which are being actively processed.""" + @abstractmethod def reproduce(self) -> Mapping[str, Mapping[str, str]]: """Reproduce queued experiments sequentially.""" diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index d0d3b65a0..a4f610520 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -8,7 +8,12 @@ from funcy import cached_property, first from dvc.exceptions import DvcException from ..exceptions import ExpQueueEmptyError -from ..executor.base import EXEC_TMP_DIR, BaseExecutor, ExecutorResult +from ..executor.base import ( + EXEC_PID_DIR, + EXEC_TMP_DIR, + BaseExecutor, + ExecutorResult, +) from ..executor.local import WorkspaceExecutor from ..refs import EXEC_BRANCH from .base import BaseStashQueue, QueueEntry, QueueGetResult @@ -17,6 +22,7 @@ from .tasks import setup_exp if TYPE_CHECKING: from dvc.repo.experiments import Experiments from dvc_task.app import FSApp + from dvc_task.proc.manager import ProcessManager from dvc_task.worker import TemporaryWorker logger = logging.getLogger(__name__) @@ -49,6 +55,13 @@ class LocalCeleryQueue(BaseStashQueue): ) return app + @cached_property + def proc(self) -> "ProcessManager": + from dvc_task.proc.manager import ProcessManager + + pid_dir = os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR) + return ProcessManager(pid_dir) + @cached_property def worker(self) -> "TemporaryWorker": from dvc_task.worker import TemporaryWorker @@ -86,6 +99,20 @@ class LocalCeleryQueue(BaseStashQueue): entry_dict = kwargs.get("entry_dict", args[0]) yield QueueEntry.from_dict(entry_dict) + def _iter_processed(self) -> Generator[QueueEntry, None, None]: + for msg in self.celery.iter_processed(): + if msg.headers.get("task") != setup_exp.name: + continue + args, kwargs, _embed = msg.decode() + entry_dict = kwargs.get("entry_dict", args[0]) + yield QueueEntry.from_dict(entry_dict) + + def iter_active(self) -> Generator[QueueEntry, None, None]: + for entry in self._iter_processed(): + proc_info = self.proc.get(entry.stash_rev) + if proc_info is not None and proc_info.returncode is None: + yield entry + def reproduce(self) -> Mapping[str, Mapping[str, str]]: raise NotImplementedError @@ -123,6 +150,11 @@ class WorkspaceQueue(BaseStashQueue): entry.name, ) + def iter_active(self) -> Generator[QueueEntry, None, None]: + # Workspace run state is reflected in the workspace itself and does not + # need to be handled via the queue + raise NotImplementedError + def reproduce(self) -> Dict[str, Dict[str, str]]: results: Dict[str, Dict[str, str]] = defaultdict(dict) try: diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index e04d4f526..7ba59101b 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -7,7 +7,7 @@ from celery.utils.log import get_task_logger from dvc.utils.fs import makedirs, remove -from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR, ExecutorInfo +from ..executor.base import ExecutorInfo from ..executor.local import TempDirExecutor from .base import BaseStashQueue, QueueEntry @@ -26,15 +26,8 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None: executor = BaseStashQueue.setup_executor( repo.experiments, entry, TempDirExecutor ) - pid_dir = os.path.join( - repo.tmp_dir, - EXEC_TMP_DIR, - EXEC_PID_DIR, - ) - infofile = os.path.join( - pid_dir, - f"{entry.stash_rev}{TempDirExecutor.INFOFILE_EXT}", - ) + 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) @@ -42,7 +35,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None: # schedule execution + cleanup chain( - run.si(cmd, wdir=pid_dir, name=entry.stash_rev), + proc.run(cmd, name=entry.stash_rev), collect_exp.s(entry.asdict(), infofile), cleanup_exp.si(entry.asdict(), executor.root_dir), ).delay()