Add join_threads() (#10086)

* Update __init__.py

Signed-off-by: Glenn Jocher <[email protected]>

* Update __init__.py

Signed-off-by: Glenn Jocher <[email protected]>

* Update __init__.py

Signed-off-by: Glenn Jocher <[email protected]>

Signed-off-by: Glenn Jocher <[email protected]>

Author: glenn-jocher

utils/__init__.py

@@ -37,6 +37,16 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def join_threads(verbose=False):
+    # Join all daemon threads, i.e. atexit.register(lambda: join_threads())
+    main_thread = threading.current_thread()
+    for t in threading.enumerate():
+        if t is not main_thread:
+            if verbose:
+                print(f'Joining thread {t.name}')
+            t.join()
+
+
 def notebook_init(verbose=True):
     # Check system software and hardware
     print('Checking setup...')