exp: implement celery exp task collection

dvc/repo/experiments/queue/base.py

@@ -18,7 +18,13 @@ from dvc.dependency.param import MissingParamsError
 from dvc.env import DVCLIVE_RESUME
 from dvc.exceptions import DvcException
 
-from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR, BaseExecutor
+from ..exceptions import CheckpointExistsError, ExperimentExistsError
+from ..executor.base import (
+    EXEC_PID_DIR,
+    EXEC_TMP_DIR,
+    BaseExecutor,
+    ExecutorResult,
+)
 from ..executor.local import WorkspaceExecutor
 from ..refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE, ExpRefInfo
 from ..stash import ExpStash, ExpStashEntry
@@ -105,6 +111,10 @@ class BaseStashQueue(ABC):
     def get(self) -> QueueGetResult:
         """Pop and return the first item in the queue."""
 
+    @abstractmethod
+    def reproduce(self) -> Mapping[str, Mapping[str, str]]:
+        """Reproduce queued experiments sequentially."""
+
     def _stash_exp(
         self,
         *args,
@@ -392,6 +402,33 @@ class BaseStashQueue(ABC):
             f"{name}{BaseExecutor.INFOFILE_EXT}",
         )
 
-    @abstractmethod
-    def reproduce(self) -> Mapping[str, Mapping[str, str]]:
-        """Reproduce queued experiments sequentially."""
+    @staticmethod
+    @scm_locked
+    def collect_executor(
+        exp: "Experiments",
+        executor: BaseExecutor,
+        exec_result: ExecutorResult,
+    ) -> Dict[str, str]:
+        results = {}
+
+        def on_diverged(ref: str, checkpoint: bool):
+            ref_info = ExpRefInfo.from_ref(ref)
+            if checkpoint:
+                raise CheckpointExistsError(ref_info.name)
+            raise ExperimentExistsError(ref_info.name)
+
+        for ref in executor.fetch_exps(
+            exp.scm,
+            force=exec_result.force,
+            on_diverged=on_diverged,
+        ):
+            exp_rev = exp.scm.get_ref(ref)
+            if exp_rev:
+                assert exec_result.exp_hash
+                logger.debug("Collected experiment '%s'.", exp_rev[:7])
+                results[exp_rev] = exec_result.exp_hash
+
+        if exec_result.ref_info is not None:
+            executor.collect_cache(exp.repo, exec_result.ref_info)
+
+        return results

dvc/repo/experiments/queue/local.py

@@ -8,15 +8,15 @@ from funcy import cached_property, first
 from dvc.exceptions import DvcException
 
 from ..exceptions import ExpQueueEmptyError
-from ..executor.base import EXEC_TMP_DIR, BaseExecutor
+from ..executor.base import EXEC_TMP_DIR, BaseExecutor, ExecutorResult
 from ..executor.local import WorkspaceExecutor
 from ..refs import EXEC_BRANCH
 from .base import BaseStashQueue, QueueEntry, QueueGetResult
 from .tasks import setup_exp
 
 if TYPE_CHECKING:
-    from celery import Celery
-
+    from dvc.repo.experiments import Experiments
+    from dvc_task.app import FSApp
     from dvc_task.worker import TemporaryWorker
 
 logger = logging.getLogger(__name__)
@@ -35,31 +35,42 @@ class LocalCeleryQueue(BaseStashQueue):
         return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, self.CELERY_DIR)
 
     @cached_property
-    def celery(self) -> "Celery":
-        from celery import Celery
+    def celery(self) -> "FSApp":
+        from dvc_task.app import FSApp
 
-        from .celery import get_config
-
-        app = Celery()
-        app.conf.update(get_config(self.wdir))
+        app = FSApp(
+            "dvc-exp-local",
+            wdir=self.wdir,
+            mkdir=True,
+            include=[
+                "dvc.repo.experiments.queue.tasks",
+                "dvc_task.proc.tasks",
+            ],
+        )
         return app
 
     @cached_property
     def worker(self) -> "TemporaryWorker":
         from dvc_task.worker import TemporaryWorker
 
-        return TemporaryWorker(self.celery, concurrency=1)
+        return TemporaryWorker(self.celery, concurrency=1, timeout=10)
+
+    def spawn_worker(self):
+        from dvc_task.proc.process import ManagedProcess
+
+        logger.debug("Spawning exp queue worker")
+        ManagedProcess.spawn(
+            ["dvc", "exp", "queue-worker", "dvc-exp1@localhost"],
+            wdir=self.wdir,
+            name="dvc-exp-worker",
+        )
 
     def put(self, *args, **kwargs) -> QueueEntry:
         """Stash an experiment and add it to the queue."""
-        from celery import chain
-
-        entry = super()._stash_exp(*args, **kwargs)
+        entry = self._stash_exp(*args, **kwargs)
         # schedule executor setup
         # TODO: chain separated git/dvc setup tasks
-        chain(
-            self.celery.tasks[setup_exp.name](entry.asdict()),
-        ).delay()
+        self.celery.tasks[setup_exp.name].delay(entry.asdict())
         return entry
 
     # NOTE: Queue consumption should not be done directly. Celery worker(s)
@@ -125,7 +136,9 @@ class WorkspaceQueue(BaseStashQueue):
                 )
             if exec_result.ref_info:
                 results[rev].update(
-                    self._collect_executor(executor, exec_result)
+                    self.collect_executor(
+                        self.repo.experiments, executor, exec_result
+                    )
                 )
         except CheckpointKilledError:
             # Checkpoint errors have already been logged
@@ -140,12 +153,16 @@ class WorkspaceQueue(BaseStashQueue):
             executor.cleanup()
         return results
 
-    def _collect_executor(  # pylint: disable=unused-argument
-        self, executor, exec_result
+    @staticmethod
+    def collect_executor(  # pylint: disable=unused-argument
+        exp: "Experiments",
+        executor: BaseExecutor,
+        exec_result: ExecutorResult,
     ) -> Dict[str, str]:
-        results = {}
-        exp_rev = self.scm.get_ref(EXEC_BRANCH)
+        results: Dict[str, str] = {}
+        exp_rev = exp.scm.get_ref(EXEC_BRANCH)
         if exp_rev:
+            assert exec_result.exp_hash
             logger.debug("Collected experiment '%s'.", exp_rev[:7])
             results[exp_rev] = exec_result.exp_hash
         return results

dvc/repo/experiments/queue/tasks.py

@@ -3,13 +3,16 @@ 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 remove
+from dvc.utils.fs import makedirs, remove
 
-from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR
+from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR, ExecutorInfo
 from ..executor.local import TempDirExecutor
 from .base import BaseStashQueue, QueueEntry
 
+logger = get_task_logger(__name__)
+
 
 @shared_task
 def setup_exp(entry_dict: Dict[str, Any]) -> None:
@@ -23,22 +26,61 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None:
     executor = BaseStashQueue.setup_executor(
         repo.experiments, entry, TempDirExecutor
     )
-    infofile = os.path.join(repo.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR)
+    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}",
+    )
+    makedirs(os.path.dirname(infofile), exist_ok=True)
     with open(infofile, "w", encoding="utf-8") as fobj:
         json.dump(executor.info.asdict(), fobj)
     cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
 
     # schedule execution + cleanup
     chain(
-        run(cmd),  # pylint: disable=no-value-for-parameter
-        cleanup_exp(  # pylint: disable=no-value-for-parameter
-            executor.root_dir, entry.asdict()
-        ),
+        run.si(cmd, wdir=pid_dir, name=entry.stash_rev),
+        collect_exp.s(entry.asdict(), infofile),
+        cleanup_exp.si(entry.asdict(), executor.root_dir),
     ).delay()
 
 
 @shared_task
+def collect_exp(
+    proc_dict: Dict[str, Any],
+    entry_dict: Dict[str, Any],
+    infofile: str,
+) -> None:
+    from dvc.repo import Repo
+    from dvc_task.proc.process import ProcessInfo
+
+    proc_info = ProcessInfo.from_dict(proc_dict)
+    if proc_info.returncode != 0:
+        # TODO: handle errors, track failed exps separately
+        pass
+
+    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))
+    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")
+
+
+@shared_task(ignore_result=True)
 def cleanup_exp(  # pylint: disable=unused-argument
-    proc_dict: Dict[str, Any], entry: Dict[str, Any], tmp_dir: str
+    entry_dict: Dict[str, Any], tmp_dir: str
 ) -> None:
     remove(tmp_dir)