_shutdown_pool now runs pool.terminate() in a daemon thread with a _TERMINATE_GRACE_SECONDS wait

Author: A9isha

src/maxtext/trainers/post_train/rl/math_verify_pool.py

@@ -28,6 +28,7 @@
 import itertools
 import multiprocessing
 import os
+import threading
 import time
 import uuid
 from etils import epath
@@ -155,8 +156,16 @@ def _shutdown_pool():
       if w.is_alive():
         w.kill()
         w.join(timeout=1.0)
-    pool.terminate()
-    pool.join()
+    # Workers SIGKILLed mid-write leave the outqueue lock orphaned, so
+    # pool.terminate() / pool.join() block forever on the internal
+    # _result_handler / _task_handler threads. Run terminate in a daemon
+    # thread with a bounded wait: pool._state flips to TERMINATE so the
+    # worker-handler stops spawning replacements, and we return even if the
+    # handler threads never unblock. Those threads leak, but they are daemon
+    # and cheap; a stuck trainer is not.
+    t = threading.Thread(target=pool.terminate, daemon=True)
+    t.start()
+    t.join(timeout=_TERMINATE_GRACE_SECONDS)
   except Exception:
     pass
 
@@ -176,7 +185,7 @@ def _get_pool(num_procs):
   return _POOL
 
 
-# ensures global worker pool is cleanly shut down when progam finishes execution
+# ensures global worker pool is cleanly shut down when program finishes execution
 atexit.register(_shutdown_pool)