Fix #316: SIGSEV during long running training when using asynchronous...

torchft/futures.py

@@ -28,11 +28,18 @@ class _TimerHandle:
     def __init__(self) -> None:
         self._lock = threading.Lock()
         self._timer_handle: Optional[asyncio.TimerHandle] = None
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._cancelled = False
 
-    def set_timer_handle(self, timer_handle: asyncio.TimerHandle) -> None:
+    def set_timer_handle(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        timer_handle: asyncio.TimerHandle,
+    ) -> None:
         with self._lock:
+            self._loop = loop
             if self._cancelled:
+                # Called from within the event loop thread, so direct cancel is safe.
                 timer_handle.cancel()
                 self._timer_handle = None
             else:
@@ -43,8 +50,16 @@ def cancel(self) -> None:
             assert not self._cancelled, "timer can only be cancelled once"
             self._cancelled = True
             if self._timer_handle is not None:
-                self._timer_handle.cancel()
+                timer_handle = self._timer_handle
+                loop = self._loop
                 self._timer_handle = None
+                if loop is not None:
+                    # Schedule cancellation on the event loop thread to avoid
+                    # SIGSEGV when using uvloop, where uv_timer_stop is not
+                    # thread-safe and must be called from the event loop thread.
+                    loop.call_soon_threadsafe(timer_handle.cancel)
+                else:
+                    timer_handle.cancel()
 
 
 class _TimeoutManager:
@@ -228,7 +243,7 @@ def _register_callback(
             timeout.total_seconds(),
             callback,
         )
-        handle.set_timer_handle(timer_handle)
+        handle.set_timer_handle(loop, timer_handle)
 
     @contextmanager
     def context_timeout(

torchft/futures_test.py

@@ -13,6 +13,7 @@
 from torch.futures import Future
 from torchft.futures import (
     _TIMEOUT_MANAGER,
+    _TimerHandle,
     context_timeout,
     future_timeout,
     future_wait,
@@ -97,6 +98,23 @@ def callback() -> None:
 
         self.assertEqual(_TIMEOUT_MANAGER._clear_del_queue(), 1)
 
+    def test_timer_handle_cancel_uses_call_soon_threadsafe(self) -> None:
+        """Cancelling from outside the event loop must go via call_soon_threadsafe.
+
+        With uvloop the underlying uv_timer_stop is NOT thread-safe; calling it
+        directly from a non-event-loop thread causes a SIGSEGV (issue #316).
+        """
+        mock_loop = Mock()
+        mock_timer_handle = Mock()
+
+        handle = _TimerHandle()
+        handle.set_timer_handle(mock_loop, mock_timer_handle)
+        handle.cancel()
+
+        # cancel() must delegate to call_soon_threadsafe, not call directly
+        mock_loop.call_soon_threadsafe.assert_called_once_with(mock_timer_handle.cancel)
+        mock_timer_handle.cancel.assert_not_called()
+
     # Test that when a timeout handle gets stuck, `sys.exit(1)` is called
     @patch("sys.exit")
     def test_exit_on_stuck_callback(self, mock_exit: Mock) -> None: