fix: make AlltoAll watchdog stop terminal

Author: chienchunhung

tensorrt_llm/_torch/alltoall_watchdog.py

@@ -197,6 +197,7 @@ def __init__(
 
         self._cv = threading.Condition()
         self._queue: Deque[_CollectiveWatch] = deque()
+        self._closed = False
         self._stopping = False
         self._thread: threading.Thread | None = None
         self._last_error: BaseException | None = None
@@ -249,6 +250,8 @@ def last_error(self) -> BaseException | None:
     def start(self) -> None:
         """Start the background polling thread. Idempotent."""
         with self._cv:
+            if self._closed:
+                raise RuntimeError("cannot start a stopped AlltoAllWatchdog")
             if self._thread is not None and self._thread.is_alive():
                 return
             self._stopping = False
@@ -262,6 +265,7 @@ def start(self) -> None:
     def stop(self, timeout_s: float | None = None) -> None:
         """Stop the polling thread and wait for it to exit."""
         with self._cv:
+            self._closed = True
             self._stopping = True
             self._queue.clear()
             self._cv.notify_all()
@@ -291,7 +295,7 @@ def watch(
 
         self.start()
         with self._cv:
-            if self._stopping:
+            if self._closed:
                 raise RuntimeError("cannot queue a stopped AlltoAllWatchdog")
             self._queue.append(
                 _CollectiveWatch(

tests/unittest/_torch/modules/test_alltoall_watchdog.py

@@ -108,6 +108,24 @@ def test_watchdog_defaults_match_design_doc() -> None:
     assert watchdog._poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S
 
 
+def test_watchdog_stop_is_terminal() -> None:
+    reader = FakeCompletionFlagReader(ep_size=1)
+    watchdog = AlltoAllWatchdog(
+        ep_size=1,
+        ep_rank=0,
+        completion_reader=reader,
+        timeout_s=0.2,
+        poll_interval_s=0.005,
+    )
+    watchdog.start()
+    watchdog.stop(timeout_s=1.0)
+
+    with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"):
+        watchdog.start()
+    with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"):
+        watchdog.watch(phase="dispatch", expected_flag=1)
+
+
 def test_wide_ep_ft_options_create_shared_health_when_enabled(monkeypatch) -> None:
     monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1")
     model_config = SimpleNamespace(