From 8dfcb0a36d002249f7435e914f25f12967b4a9dc Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Thu, 9 Jun 2022 13:46:03 +0900 Subject: [PATCH] exp run: update Ctrl+C handling for --temp and --run-all --- dvc/repo/experiments/__init__.py | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 39c4aaf4b..34c432696 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -172,19 +172,25 @@ class Experiments: # TODO: re-enable --jobs concurrency self.celery_queue.spawn_worker() failed = [] - for entry in entries: - # wait for task execution to start - while not self.celery_queue.proc.get(entry.stash_rev): - time.sleep(1) - self.celery_queue.follow(entry) - # wait for task collection to complete - result = self.celery_queue.get_result(entry) - if result is None or result.exp_hash is None: - name = entry.name or entry.stash_rev[:7] - failed.append(name) - elif result.ref_info: - exp_rev = self.scm.get_ref(str(result.ref_info)) - results[exp_rev] = result.exp_hash + try: + for entry in entries: + # wait for task execution to start + while not self.celery_queue.proc.get(entry.stash_rev): + time.sleep(1) + self.celery_queue.follow(entry) + # wait for task collection to complete + result = self.celery_queue.get_result(entry) + if result is None or result.exp_hash is None: + name = entry.name or entry.stash_rev[:7] + failed.append(name) + elif result.ref_info: + exp_rev = self.scm.get_ref(str(result.ref_info)) + results[exp_rev] = result.exp_hash + except KeyboardInterrupt: + ui.write( + "Experiment(s) are still executing in the background. To " + "abort execution use 'dvc queue kill' or 'dvc queue stop'." + ) if failed: names = ", ".join(name for name in failed) ui.error(f"Failed to reproduce experiment(s) '{names}'")