Keep workers reconnecting across hub outages

docs/environment.md

@@ -33,7 +33,6 @@ this page covers the worker itself.
 | `WORKER_ID` | `worker_id` | per-pod identity |
 | `ENDPOINT_LOCK_PATH` | `endpoint_lock_path` | discovery manifest path (baked default in images) |
 | `RUNPOD_POD_ID` | `runpod_pod_id` | set by the RunPod runtime |
-| `WORKER_DISCONNECTED_TIMEOUT_S` | `worker_disconnected_timeout_s` | self-exit window when orchestrator unreachable |
 | `WORKER_IMAGE_DIGEST` | `worker_image_digest` | provenance stamped by the image build — currently no launcher sets it (pgw#514/P4); kept for a future tensorhub stamp |
 
 ## Tuning knobs (Settings fields)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "gen-worker"
-version = "0.26.2"
+version = "0.26.3"
 description = "A library used to build custom functions in Cozy Creator's serverless function platform."
 readme = "README.md"
 license = "MIT"

src/gen_worker/config/loader.py

@@ -37,7 +37,6 @@
     "WORKER_JWT": "worker_jwt",
     "ENDPOINT_LOCK_PATH": "endpoint_lock_path",
     "RUNPOD_POD_ID": "runpod_pod_id",
-    "WORKER_DISCONNECTED_TIMEOUT_S": "worker_disconnected_timeout_s",
     "WORKER_IMAGE_DIGEST": "worker_image_digest",
     "TENSORHUB_URL": "tensorhub_url",
     "TENSORHUB_TOKEN": "tensorhub_token",

src/gen_worker/config/settings.py

@@ -51,16 +51,6 @@ class Settings(msgspec.Struct, frozen=True, kw_only=True):
     # Runtime introspection (set by the RunPod runtime; not configuration).
     runpod_pod_id: str = ""
 
-    # Worker self-exit timeout when the orchestrator is unreachable.
-    # The reconnect loop is otherwise infinite (correct for transient
-    # network blips and single-replica restarts). After this many seconds
-    # of NO successful connection, the worker exits cleanly so its
-    # container is reaped by docker/runpod — handles stack-shutdown
-    # ergonomics (compose down, machine power-off) without a separate
-    # shutdown-coordination protocol. Default 600s = 10 min. gen-orchestrator
-    # issue #317.
-    worker_disconnected_timeout_s: int = 600
-
     # Provenance stamped into WorkerResources by the image build / launcher.
     # TODO(pgw#514/P4): nothing produces WORKER_IMAGE_DIGEST today (verified
     # against gen-orchestrator/tensorhub/e2e — no launcher sets it), so this

src/gen_worker/lifecycle.py

@@ -444,9 +444,6 @@ def _begin_drain(self, deadline_ms: int) -> None:
         deadline_s = (deadline_ms / 1000.0) if deadline_ms > 0 else None
         loop = asyncio.get_running_loop()
         self._drain_deadline_at = loop.time() + deadline_s if deadline_s is not None else None
-        if self._drain_deadline_at is None and self.transport is not None:
-            self.transport.begin_non_expiring_drain()
-
     async def _finish_drain(self) -> None:
         deadline_at = self._drain_deadline_at
         loop = asyncio.get_running_loop()

src/gen_worker/transport.py

@@ -36,8 +36,7 @@
 _BACKOFF_RESET_AFTER_S = 60.0
 _HELLO_ACK_TIMEOUT_S = 30.0
 # FAILED_PRECONDITION details that can never heal by retrying: identity is
-# wrong for this deployment, so exit for reap instead of burning the full
-# disconnected timeout.
+# wrong for this deployment, so retrying cannot repair it.
 _PERMANENT_PRECONDITION_MARKERS = (
     "worker_id_mismatch",
     "release_id_mismatch",
@@ -46,8 +45,7 @@
 
 
 class FatalTransportError(Exception):
-    """Unrecoverable: protocol mismatch, repeated auth rejection, or the
-    disconnected-timeout elapsed. The worker process should exit."""
+    """Unrecoverable protocol mismatch or persistent registration rejection."""
 
 
 def normalize_grpc_addr(addr: str, default_tls: Optional[bool] = None) -> Tuple[str, bool]:
@@ -188,7 +186,6 @@ def __init__(
         self._stopping = asyncio.Event()
         self._connected = asyncio.Event()
         self._clean_close = False
-        self._non_expiring_drain = False
         self.reconnect_delays: List[float] = []  # observability + tests
         self._consecutive_auth_failures = 0
         self._first_auth_failure_at: Optional[float] = None
@@ -222,10 +219,6 @@ async def wait_connected(self, timeout: Optional[float] = None) -> bool:
 
     # ---- drain / shutdown --------------------------------------------------
 
-    def begin_non_expiring_drain(self) -> None:
-        """Keep reconnecting until a zero-deadline drain ships its results."""
-        self._non_expiring_drain = True
-
     async def close_after_flush(self, timeout: Optional[float] = None) -> bool:
         """Ship the queue, then stop; ``None`` waits until every result ships."""
         flushed = await self.queue.wait_empty(timeout=timeout)
@@ -271,14 +264,11 @@ def _metadata(self) -> Optional[List[Tuple[str, str]]]:
         return [("authorization", f"Bearer {token}")]
 
     async def run(self) -> None:
-        """Reconnect loop. Returns cleanly on drain close; raises
-        FatalTransportError on version mismatch / auth lockout / timeout."""
+        """Reconnect until stopped; fatal protocol/auth failures still exit."""
         attempt = 0
         redirect_addr: Optional[str] = None
         redirect_tls: Optional[bool] = None
         redirect_hops = 0
-        last_connected = time.monotonic()
-
         while not self._stopping.is_set():
             if redirect_addr is not None:
                 # Schemeless redirect targets inherit the TLS mode of the
@@ -346,18 +336,8 @@ async def run(self) -> None:
 
             now = time.monotonic()
             if self._connected_at is not None:
-                last_connected = now
                 if now - self._connected_at >= _BACKOFF_RESET_AFTER_S:
                     attempt = 0
-            timeout_s = float(self._settings.worker_disconnected_timeout_s or 0)
-            if (
-                timeout_s > 0
-                and now - last_connected > timeout_s
-                and not self._non_expiring_drain
-            ):
-                raise FatalTransportError(
-                    f"no successful connection for {timeout_s:.0f}s; exiting for reap"
-                )
 
             delay = random.uniform(0, min(self._backoff_cap, self._backoff_base * (2 ** attempt)))
             attempt += 1

tests/test_drain_flush.py

@@ -6,7 +6,7 @@
 from gen_worker.executor import Executor
 from gen_worker.lifecycle import Lifecycle
 from gen_worker.pb import worker_scheduler_pb2 as pb
-from gen_worker.transport import FatalTransportError, Transport
+from gen_worker.transport import Transport
 
 
 class _IdleExecutor(Executor):
@@ -57,10 +57,7 @@ async def on_disconnect(self) -> None:
 class _ReconnectTransport(Transport):
     def __init__(self) -> None:
         super().__init__(
-            Settings(
-                orchestrator_public_addr="localhost:1",
-                worker_disconnected_timeout_s=1,
-            ),
+            Settings(orchestrator_public_addr="localhost:1"),
             _DisconnectHandlers(),
             backoff_base_s=0.0,
             backoff_cap_s=0.0,
@@ -70,7 +67,6 @@ def __init__(self) -> None:
     async def _connect_once(self, target: str, use_tls: bool) -> None:
         self.connect_attempts += 1
         if self.connect_attempts == 1:
-            await asyncio.sleep(1.01)
             raise ConnectionError("orchestrator unavailable")
         kind, result = await self.queue.get()
         assert kind == "result"
@@ -84,7 +80,7 @@ def _result() -> pb.WorkerMessage:
     ))
 
 
-def test_non_expiring_drain_waits_for_blocked_result_queue() -> None:
+def test_unbounded_drain_waits_for_blocked_result_queue() -> None:
     async def _run() -> None:
         transport = Transport(Settings(orchestrator_public_addr="localhost:1"), object())
         executor = _IdleExecutor()
@@ -121,7 +117,6 @@ async def _run() -> None:
         assert lifecycle.drained.is_set()
         assert transport._stopping.is_set()
         assert not transport._clean_close
-        assert not transport._non_expiring_drain
         assert transport.queue.pending_result_keys == [("request-1", 1)]
         assert executor.wait_timeout is not None
         assert transport.flush_timeout is not None
@@ -151,7 +146,7 @@ async def _run() -> None:
     asyncio.run(_run())
 
 
-def test_non_expiring_drain_reconnects_past_disconnected_timeout() -> None:
+def test_unbounded_drain_reconnects_until_result_ships() -> None:
     async def _run() -> None:
         transport = _ReconnectTransport()
         executor = _IdleExecutor()
@@ -169,17 +164,3 @@ async def _run() -> None:
         assert transport.queue.pending_result_keys == []
 
     asyncio.run(_run())
-
-
-def test_disconnected_timeout_still_applies_without_non_expiring_drain() -> None:
-    async def _run() -> None:
-        transport = _ReconnectTransport()
-        try:
-            await transport.run()
-        except FatalTransportError as exc:
-            assert "no successful connection" in str(exc)
-        else:
-            raise AssertionError("ordinary disconnect unexpectedly retried past its lifetime")
-        assert transport.connect_attempts == 1
-
-    asyncio.run(_run())

tests/test_token_refresh.py

@@ -124,7 +124,6 @@ async def scenario() -> None:
             orchestrator_public_addr=f"127.0.0.1:{port}",
             worker_id="rotation-worker",
             worker_jwt="boot-token",
-            worker_disconnected_timeout_s=60,
         )
         handlers = _Handlers()
         transport = Transport(settings, handlers, backoff_base_s=0.05, backoff_cap_s=0.2)
@@ -181,7 +180,6 @@ async def scenario() -> None:
             orchestrator_public_addr=f"127.0.0.1:{port}",
             worker_id="rotation-worker",
             worker_jwt="boot-token",
-            worker_disconnected_timeout_s=60,
         )
         transport = Transport(settings, _Handlers(), backoff_base_s=0.05, backoff_cap_s=0.2)
         run_task = asyncio.create_task(transport.run())

tests/test_worker_grpc_e2e.py

@@ -166,7 +166,6 @@ def __init__(self, scheduler: FakeScheduler, port: int) -> None:
             orchestrator_public_addr=f"127.0.0.1:{port}",
             worker_id="e2e-worker",
             worker_jwt="",
-            worker_disconnected_timeout_s=60,
         )
         self.worker = Worker(
             settings,
@@ -336,6 +335,41 @@ def test_full_contract_round_trip(scheduler_and_worker) -> None:
     assert harness.join() == 0
 
 
+def test_worker_reconnects_after_hub_restart() -> None:
+    scheduler = FakeScheduler()
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
+    pb_grpc.add_WorkerSchedulerServicer_to_server(scheduler, server)
+    port = server.add_insecure_port("127.0.0.1:0")
+    server.start()
+    harness = _Harness(scheduler, port)
+    harness.start()
+    replacement = None
+    try:
+        scheduler.wait_connection(0)
+        before = len(harness.worker.transport.reconnect_delays)
+        assert server.stop(grace=0).wait(_TIMEOUT)
+
+        deadline = time.monotonic() + _TIMEOUT
+        while len(harness.worker.transport.reconnect_delays) < before + 4:
+            assert harness._thread.is_alive(), "worker exited while hub was down"
+            assert time.monotonic() < deadline, "worker did not keep reconnecting"
+            time.sleep(0.02)
+
+        replacement_scheduler = FakeScheduler()
+        replacement = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
+        pb_grpc.add_WorkerSchedulerServicer_to_server(replacement_scheduler, replacement)
+        assert replacement.add_insecure_port(f"127.0.0.1:{port}") == port
+        replacement.start()
+
+        assert replacement_scheduler.wait_connection(0).hello is not None
+        assert harness.exit_code is None
+    finally:
+        harness.stop()
+        server.stop(grace=0)
+        if replacement is not None:
+            replacement.stop(grace=0)
+
+
 def test_deadline_marks_fatal_and_frees_the_worker(scheduler_and_worker) -> None:
     scheduler, harness = scheduler_and_worker
     conn = scheduler.wait_connection(0)
@@ -594,7 +628,7 @@ def test_auth_rejection_within_window_keeps_retrying() -> None:
 
 def test_permanent_precondition_exits_fast() -> None:
     """worker_id_mismatch cannot heal by retrying: exit for reap immediately
-    instead of burning the disconnected timeout (#372)."""
+    instead of retrying forever (#372)."""
 
     class _Mismatch(FakeScheduler):
         def Connect(self, request_iterator, context):