Improve handling of KeyboardInterrupt in CLI applications (#567)

learning-at-home · Jun 11, 2023 · 7a91d03 · 7a91d03
1 parent 0d2614d
commit 7a91d03
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/hivemind/hivemind_cli/run_dht.py b/hivemind/hivemind_cli/run_dht.py
@@ -84,9 +84,14 @@ def main():
     )
     log_visible_maddrs(dht.get_visible_maddrs(), only_p2p=args.use_ipfs)
 
-    while True:
-        dht.run_coroutine(report_status, return_future=False)
-        time.sleep(args.refresh_period)
+    try:
+        while True:
+            dht.run_coroutine(report_status, return_future=False)
+            time.sleep(args.refresh_period)
+    except KeyboardInterrupt:
+        logger.info("Caught KeyboardInterrupt, shutting down")
+    finally:
+        dht.shutdown()
 
 
 if __name__ == "__main__":

diff --git a/hivemind/hivemind_cli/run_server.py b/hivemind/hivemind_cli/run_server.py
@@ -108,8 +108,6 @@ def main():
         server.join()
     except KeyboardInterrupt:
         logger.info("Caught KeyboardInterrupt, shutting down")
-    finally:
-        server.shutdown()
 
 
 if __name__ == "__main__":