Graceful shutdown

erikbern · Apr 16, 2023 · f048559 · f048559
1 parent 6521371
commit f048559
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/ann_benchmarks/main.py b/ann_benchmarks/main.py
@@ -208,8 +208,10 @@ def _test(df):
         raise Exception(
             f"Batch mode uses all available CPU resources, --parallelism should be set to 1. (Was: {args.parallelism})"
         )
-    workers = [multiprocessing.Process(target=run_worker, args=(i + 1, args, queue)) for i in range(args.parallelism)]
-    [worker.start() for worker in workers]
-    [worker.join() for worker in workers]
-
-    # TODO: need to figure out cleanup handling here
+    try:
+        workers = [multiprocessing.Process(target=run_worker, args=(i + 1, args, queue)) for i in range(args.parallelism)]
+        [worker.start() for worker in workers]
+        [worker.join() for worker in workers]
+    finally:
+        logger.info("Terminating %d workers" % len(workers))
+        [worker.terminate() for worker in workers]
diff --git a/ann_benchmarks/runner.py b/ann_benchmarks/runner.py
@@ -255,6 +255,7 @@ def stream_logs():
         logger.error("Container.wait for container %s failed with exception" % container.short_id)
         traceback.print_exc()
     finally:
+        logger.info("Removing container")
         container.remove(force=True)