Bound orchestrator shutdown waits

tests/experimental/orchestrator/orchestrator_test.py

@@ -14,6 +14,7 @@
 
 """Unit tests for ClusterOrchestrator."""
 
+import time
 from unittest import mock
 
 from absl.testing import absltest
@@ -133,6 +134,27 @@ def test_bring_up_and_shutdown_remote_worker_handles(self):
     mock_rollout.submit.assert_any_call("stop")
     mock_actor.submit.assert_any_call("stop")
 
+  def test_shutdown_survives_a_wedged_worker(self):
+    from tunix.experimental.worker import remote_execution
+
+    wedged = mock.MagicMock(spec=remote_execution.ActorHandle)
+    wedged.submit.side_effect = lambda *a, **kw: time.sleep(5)
+    healthy = mock.MagicMock(spec=remote_execution.ActorHandle)
+
+    registry = worker_registry.WorkerRegistry()
+    orch = orchestrator.ClusterOrchestrator(
+        registry=registry,
+        lifecycle_driver=self.mock_lifecycle,
+        monitor=self.mock_monitor,
+    )
+    orch.register_worker_handle("rollout-0", [datatypes.Role.ROLLOUT], wedged)
+    orch.register_worker_handle("actor-0", [datatypes.Role.ACTOR], healthy)
+
+    with mock.patch.object(orchestrator, "_STOP_TIMEOUT_S", 0.2):
+      orch._shutdown_remote_workers()
+
+    healthy.submit.assert_any_call("stop")
+
   def test_create_engine_wraps_local_workers_as_in_process_handles(self):
     from tunix.experimental.worker import remote_execution
 

tunix/experimental/orchestrator/orchestrator.py

@@ -21,6 +21,7 @@
 
 import collections
 from collections.abc import Callable, Iterable, Sequence
+from concurrent import futures
 from typing import Any
 
 from absl import logging
@@ -37,6 +38,7 @@
 from tunix.experimental.worker import abstract_worker
 from tunix.experimental.worker import remote_execution
 
+_STOP_TIMEOUT_S = 60.0  # Timeout for stopping remote workers. 60 should not be touched for any healthy stop.
 
 class ClusterOrchestrator:
   """Supervises cluster hardware, health monitoring, and program execution."""
@@ -192,12 +194,20 @@ def _bring_up_remote_workers(self, dummy_data: Any = None) -> None:
       self._remote_worker_handles_by_id[worker_id].submit("start")
 
   def _shutdown_remote_workers(self) -> None:
-    """Stops remote worker handles best-effort."""
-    for worker_id in sorted(self._remote_worker_infos):
+    """Stops remote worker handles best-effort, with a hard timeout."""
+    pool = futures.ThreadPoolExecutor(max_workers=4)
+    stops = {
+        worker_id: pool.submit(
+            self._remote_worker_handles_by_id[worker_id].submit, "stop"
+        )
+        for worker_id in sorted(self._remote_worker_infos)
+    }
+    for worker_id, fut in stops.items():
       try:
-        self._remote_worker_handles_by_id[worker_id].submit("stop")
+        fut.result(timeout=_STOP_TIMEOUT_S)
       except Exception as err:  # pylint: disable=broad-except
         logging.warning("Failed to stop remote worker %s: %r", worker_id, err)
+    pool.shutdown(wait=False)
 
   def _create_engine(self) -> distributed_rl_engine.DistributedRLEngine:
     """Constructs a DistributedRLEngine from the registered role groups."""