From ccbe2de89eec604838fd4ae786751f3baa06ad1d Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Tue, 28 Feb 2023 17:49:25 +0900 Subject: [PATCH] exp: auto cleanup celery on queue worker exit --- dvc/commands/experiments/queue_worker.py | 9 ++++++++- dvc/repo/experiments/queue/celery.py | 5 ++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/dvc/commands/experiments/queue_worker.py b/dvc/commands/experiments/queue_worker.py index 81b6372d9..b85ff1d7a 100644 --- a/dvc/commands/experiments/queue_worker.py +++ b/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) diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index e85b1cf14..8d5364429 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/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)