exp: use a single chain for all exp tasks

- removes separated scheduling for setup_exp

dvc/repo/experiments/__init__.py

@@ -360,7 +360,6 @@ class Experiments:
             return exp_baseline
         raise BaselineMismatchError(exp_baseline, baseline_sha)
 
-    @scm_locked
     def get_baseline(self, rev):
         """Return the baseline rev for an experiment rev."""
         return self._get_baseline(rev)

dvc/repo/experiments/queue/base.py

@@ -190,15 +190,8 @@ class BaseStashQueue(ABC):
         """Iterate over items in the queue."""
 
     @abstractmethod
-    def iter_active(
-        self, ignore_collected: bool = False
-    ) -> Generator[QueueEntry, None, None]:
-        """Iterate over items which are being actively processed.
-
-        Arguments:
-            ignore_collected: Ignore experiments which have been collected (but
-                not cleaned up).
-        """
+    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]]:

dvc/repo/experiments/queue/local.py

@@ -15,6 +15,8 @@ from typing import (
     Set,
 )
 
+from celery import chain
+from celery.canvas import Signature
 from funcy import cached_property, first
 from kombu.message import Message
 
@@ -33,7 +35,7 @@ from ..executor.local import WorkspaceExecutor
 from ..refs import EXEC_BRANCH
 from ..stash import ExpStashEntry
 from .base import BaseStashQueue, QueueEntry, QueueGetResult
-from .tasks import setup_exp
+from .tasks import cleanup_exp, collect_exp, setup_exp
 
 if TYPE_CHECKING:
     from dvc.repo.experiments import Experiments
@@ -134,11 +136,26 @@ class LocalCeleryQueue(BaseStashQueue):
     def put(self, *args, **kwargs) -> QueueEntry:
         """Stash an experiment and add it to the queue."""
         entry = self._stash_exp(*args, **kwargs)
-        # schedule executor setup
-        # TODO: chain separated git/dvc setup tasks
-        self.celery.tasks[setup_exp.name].delay(entry.asdict())
+        task_chain = self._chain_exp(entry)
+        task_chain.freeze()
+        task_chain.delay()
         return entry
 
+    def _chain_exp(self, entry: QueueEntry) -> Signature:
+        """Return a complete experiment chain for this queue entry."""
+        infofile = self.get_infofile_path(entry.stash_rev)
+        cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
+        return self.celery.signature(
+            chain(
+                setup_exp.s(entry.asdict()),
+                self.proc.run_signature(
+                    cmd, name=entry.stash_rev, immutable=True
+                ),
+                collect_exp.s(entry.asdict()),
+                cleanup_exp.s(entry.asdict()),
+            )
+        )
+
     # NOTE: Queue consumption should not be done directly. Celery worker(s)
     # will automatically consume available experiments.
     def get(self) -> QueueGetResult:
@@ -319,9 +336,7 @@ class WorkspaceQueue(BaseStashQueue):
                 entry.name,
             )
 
-    def iter_active(
-        self, ignore_collected: bool = False
-    ) -> Generator[QueueEntry, None, None]:
+    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

dvc/repo/experiments/queue/tasks.py

@@ -1,6 +1,7 @@
-from typing import Any, Dict
+from typing import Any, Dict, List
 
-from celery import chain, shared_task
+from celery import shared_task
+from celery.signals import task_postrun
 from celery.utils.log import get_task_logger
 
 from dvc.utils.fs import remove
@@ -13,18 +14,13 @@ logger = get_task_logger(__name__)
 
 
 @shared_task
-def setup_exp(entry_dict: Dict[str, Any]) -> str:
-    """Setup (queue) an experiment.
+def setup_exp(entry_dict: Dict[str, Any]) -> None:
+    """Setup an experiment.
 
     Arguments:
         entry_dict: Serialized QueueEntry for this experiment.
-
-    Returns:
-        Celery task ID for the final task in the queued task chain (experiment
-        cleanup task ID).
     """
     from dvc.repo import Repo
-    from dvc_task.proc.tasks import run
 
     entry = QueueEntry.from_dict(entry_dict)
     repo = Repo(entry.dvc_root)
@@ -33,28 +29,24 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str:
     executor = BaseStashQueue.setup_executor(
         repo.experiments, entry, TempDirExecutor
     )
-    proc = repo.experiments.celery_queue.proc
     infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev)
     executor.info.dump_json(infofile)
-    cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
-
-    # schedule execution + cleanup
-    exp_chain = chain(
-        proc.run(cmd, name=entry.stash_rev),
-        collect_exp.s(entry.asdict(), infofile),
-        cleanup_exp.si(entry.asdict(), executor.root_dir),
-    )
-    cleanup_task = exp_chain.freeze()
-    exp_chain.delay()
-    return cleanup_task.id
 
 
 @shared_task
 def collect_exp(
     proc_dict: Dict[str, Any],
     entry_dict: Dict[str, Any],
-    infofile: str,
-) -> None:
+) -> str:
+    """Collect results for an experiment.
+
+    Arguments:
+        proc_dict: Serialized ProcessInfo for experiment executor process.
+        entry_dict: Serialized QueueEntry for this experiment.
+
+    Returns:
+        Directory to be cleaned up after this experiment.
+    """
     from dvc.repo import Repo
     from dvc_task.proc.process import ProcessInfo
 
@@ -65,6 +57,7 @@ def collect_exp(
 
     entry = QueueEntry.from_dict(entry_dict)
     repo = Repo(entry.dvc_root)
+    infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev)
     executor_info = ExecutorInfo.load_json(infofile)
     logger.debug("Collecting experiment info '%s'", str(executor_info))
     executor = TempDirExecutor.from_info(executor_info)
@@ -78,13 +71,38 @@ def collect_exp(
                 logger.debug("Collected experiment '%s'", rev[:7])
         else:
             logger.debug("Exec result was None")
-    finally:
-        executor_info.collected = True
-        executor_info.dump_json(infofile)
+    except Exception:  # pylint: disable=broad-except
+        # Log exceptions but do not re-raise so that task chain execution
+        # continues
+        logger.exception("Failed to collect experiment")
+    return executor.root_dir
 
 
 @shared_task
 def cleanup_exp(  # pylint: disable=unused-argument
-    entry_dict: Dict[str, Any], tmp_dir: str
+    tmp_dir: str, entry_dict: Dict[str, Any]
 ) -> None:
+    """Cleanup after an experiment.
+
+    Arguments:
+        tmp_dir: Temp directory to be removed.
+        entry_dict: Serialized QueueEntry for this experiment.
+    """
     remove(tmp_dir)
+
+
+@task_postrun.connect(sender=cleanup_exp)
+def _cleanup_postrun_handler(
+    args: List[Any] = None,
+    **kwargs,
+):
+    from dvc.repo import Repo
+
+    assert args
+    (_, entry_dict) = args
+    entry = QueueEntry.from_dict(entry_dict)
+    repo = Repo(entry.dvc_root)
+    infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev)
+    executor_info = ExecutorInfo.load_json(infofile)
+    executor_info.collected = True
+    executor_info.dump_json(infofile)

dvc/repo/experiments/show.py

@@ -195,7 +195,7 @@ def show(
             )
         # collect celery experiments
         for entry in chain(
-            repo.experiments.celery_queue.iter_active(ignore_collected=True),
+            repo.experiments.celery_queue.iter_active(),
             repo.experiments.celery_queue.iter_queued(),
         ):
             stash_rev = entry.stash_rev