exp: auto cleanup celery on queue worker exit

dvc/commands/experiments/queue_worker.py

@@ -9,7 +9,9 @@ class CmdQueueWorker(CmdBase):
     """Run the exp queue worker."""
 
     def run(self):
-        self.repo.experiments.celery_queue.worker.start(self.args.name)
+        self.repo.experiments.celery_queue.worker.start(
+            self.args.name, fsapp_clean=self.args.clean
+        )
         return 0
 
 
@@ -22,4 +24,9 @@ def add_parser(experiments_subparsers, parent_parser):
         add_help=False,
     )
     parser.add_argument("name", help="Celery worker name.")
+    parser.add_argument(
+        "--clean",
+        action="store_true",
+        help="Automatically cleanup celery broker on shutdown.",
+    )
     parser.set_defaults(func=CmdQueueWorker)

dvc/repo/experiments/queue/celery.py

@@ -90,7 +90,7 @@ class LocalCeleryQueue(BaseStashQueue):
                 "dvc_task.proc.tasks",
             ],
         )
-        app.conf.update({"task_acks_late": True})
+        app.conf.update({"task_acks_late": True, "result_expires": None})
         return app
 
     @cached_property
@@ -138,6 +138,9 @@ class LocalCeleryQueue(BaseStashQueue):
         wdir_hash = hashlib.sha256(self.wdir.encode("utf-8")).hexdigest()[:6]
         node_name = f"dvc-exp-{wdir_hash}-{num}@localhost"
         cmd = ["exp", "queue-worker", node_name]
+        if num == 1:
+            # automatically run celery cleanup when primary worker shuts down
+            cmd.append("--clean")
         name = f"dvc-exp-worker-{num}"
 
         logger.debug("start a new worker: %s, node: %s", name, node_name)