digitalkin 1.0.0.dev18__tar.gz → 1.0.0.dev20__tar.gz

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 1.0.0.dev18
+Version: 1.0.0.dev20
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev20}/pyproject.toml

@@ -37,7 +37,7 @@
         "pydantic-settings>=2.14.1",
         "redis[hiredis]>=7.4.0,<9",
     ]
-    version = "1.0.0.dev18"
+    version = "1.0.0.dev20"
 
     [project.optional-dependencies]
         performance = [ "uvloop>=0.21" ]

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev20}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "1.0.0.dev18"
+    __version__ = "1.0.0.dev20"

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev20}/src/digitalkin/core/task_manager/module_runner.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any
 
 from google.protobuf import json_format, struct_pb2
 from pydantic import ValidationError
+from redis.exceptions import RedisError
 
 from digitalkin.core.exceptions import BackpressureTimeoutError
 from digitalkin.core.profiling.step_timer import StepTimer
@@ -41,7 +42,7 @@ class ModuleRunner:
         self._redis_client = redis_client
         self._servicer = servicer
 
-    async def run(  # noqa: PLR0914, PLR0915
+    async def run(  # noqa: C901, PLR0914, PLR0915
         self,
         query: struct_pb2.Struct,
         *,
@@ -215,6 +216,16 @@ class ModuleRunner:
         except BackpressureTimeoutError as exc:
             logger.exception("ModuleRunner: backpressure timeout", extra=log_extra)
             await on_fatal(StreamErrorCode.BACKPRESSURE_TIMEOUT.value, str(exc))
+        except RedisError as exc:
+            logger.error(
+                "[VALIDATE R2] ModuleRunner: Redis unavailable writing output: task_id=%s",
+                task_id,
+                extra=log_extra,
+            )  # TODO(validate): remove after prod validation
+            await on_fatal(
+                StreamErrorCode.REDIS_UNAVAILABLE.value,
+                f"redis unavailable: {type(exc).__name__}: {exc}",
+            )
         except Exception as exc:
             logger.exception("ModuleRunner: module job failed", extra=log_extra)
             await on_fatal(

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev20}/src/digitalkin/grpc_servers/gateway_servicer.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import contextlib
 import json
+import random
 import time
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any
@@ -34,7 +35,7 @@ from digitalkin.services.communication.exceptions import InvalidConsumerAddressE
 from digitalkin.services.communication.grpc_communication import GrpcCommunication
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, AsyncIterator
+    from collections.abc import AsyncGenerator, AsyncIterator, Callable
 
     from digitalkin.core.task_manager.module_runner import ModuleRunner
     from digitalkin.core.task_manager.redis.redis_client import RedisClient
@@ -208,61 +209,36 @@ class GatewayServicer:
             return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
 
         if self._registry.get(task_id) is not None:
-            logger.debug("Dedup: session exists, reusing", extra=log_extra)
-            return gateway_pb2.StartStreamResponse(accepted=True, task_id=task_id)
+            logger.info(
+                "[VALIDATE RD1] StartStream refused: a dial is already live for task_id=%s",
+                task_id,
+                extra=log_extra,
+            )  # TODO(validate): remove after prod validation
+            return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
         timer.mark("dedup_check")
 
         # Durable at-most-once guard: survives session teardown and spans replicas.
-        # A re-issued StartStream for an already-run task is accepted but does NOT
-        # re-dial/re-execute — the caller resumes the existing output via Stream/from_seq.
-        claim = await self._idempotency.claim(task_id, SharedRedisListener.PROCESS_ID)
+        # Reconnection is server-driven (the live dial-back auto re-dials the same
+        # consumer), so a re-issued StartStream for an already-claimed/running task is
+        # REFUSED — one task_id maps to exactly one dial.
+        try:
+            claim = await self._idempotency.claim(task_id, SharedRedisListener.PROCESS_ID)
+        except RedisError:
+            logger.warning(
+                "[VALIDATE R2] StartStream claim Redis failure → accepted=False",
+                extra=log_extra,
+                exc_info=True,
+            )  # TODO(validate): remove after prod validation
+            return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
         timer.mark("idempotency_claim")
         if claim is not ClaimResult.CLAIMED:
-            # Re-issued StartStream for an already-claimed task. The prior dial-back
-            # is dead (the dedup branch above returned if its session were still live),
-            # so re-dial in resume mode to re-attach the restarted consumer from its
-            # cursor — as long as the durable stream is still retained.
-            stream_len = await self._redis_client.xlen(f"task:{task_id}:stream")
-            if stream_len == 0:
-                logger.info(
-                    "[VALIDATE IDEM1] no live stream to resume (TTL expired) task_id=%s (%s)",
-                    task_id,
-                    claim.name,
-                    extra=log_extra,
-                )  # TODO(validate): remove after prod validation
-                return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
-            # Atomic guard: recheck + register + spawn with no await that yields
-            # between them, so a concurrent re-poke falls into the dedup branch and acks.
-            if self._registry.get(task_id) is not None:
-                logger.debug("resume-dial already scheduled, reusing", extra=log_extra)
-                return gateway_pb2.StartStreamResponse(accepted=True, task_id=task_id)
-            session = StreamSession(task_id=task_id)
-            registered = await self._registry.register(
-                session,
-                setup_id=request.setup_id,
-                mission_id=request.mission_id,
-            )
-            if not registered:
-                logger.warning("Resume-dial rejected (capacity)", extra=log_extra)
-                return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
             logger.info(
-                "[VALIDATE IDEM1] resume-dial scheduled task_id=%s (%s) stream_len=%d",
+                "[VALIDATE RD1] StartStream refused: task already claimed/running task_id=%s (%s)",
                 task_id,
                 claim.name,
-                stream_len,
                 extra=log_extra,
             )  # TODO(validate): remove after prod validation
-            self._spawn(
-                self._dial_consumer(
-                    task_id=task_id,
-                    mission_id=request.mission_id,
-                    setup_id=request.setup_id,
-                    address=client_address,
-                    resume=True,
-                ),
-                name=f"dial_consumer_{task_id}",
-            )
-            return gateway_pb2.StartStreamResponse(accepted=True, task_id=task_id)
+            return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
 
         session = StreamSession(task_id=task_id)
         accepted = await self._registry.register(
@@ -274,7 +250,8 @@ class GatewayServicer:
         if not accepted:
             logger.warning("Session rejected (capacity)", extra=log_extra)
             # Release the claim so this task can be retried once capacity frees up.
-            await self._idempotency.release(task_id)
+            with contextlib.suppress(RedisError):
+                await self._idempotency.release(task_id)
             return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
 
         # Seed stream.start so the consumer's first XREAD finds data immediately.
@@ -289,10 +266,22 @@ class GatewayServicer:
             },
         })
         timer.mark("build_start_info")
-        await self._redis_client.xadd(
-            f"task:{task_id}:stream",
-            {"pb": start_info.SerializeToString(), "seq": "0"},
-        )
+        try:
+            await self._redis_client.xadd(
+                f"task:{task_id}:stream",
+                {"pb": start_info.SerializeToString(), "seq": "0"},
+            )
+        except RedisError:
+            logger.warning(
+                "[VALIDATE R2] StartStream xadd(stream.start) Redis failure → accepted=False",
+                extra=log_extra,
+                exc_info=True,
+            )  # TODO(validate): remove after prod validation
+            # Claimed + registered but the stream couldn't be seeded: undo both so a retry re-runs.
+            with contextlib.suppress(RedisError):
+                await self._idempotency.release(task_id)
+            await self._registry.unregister(task_id)
+            return gateway_pb2.StartStreamResponse(accepted=False, task_id=task_id)
         timer.mark("xadd_stream_start")
 
         logger.info("→ Dial-back scheduled to consumer %s", client_address, extra=log_extra)
@@ -424,7 +413,14 @@ class GatewayServicer:
 
         # Late client: session finished but data still in Redis.
         if session is None:
-            stream_len = await self._redis_client.xlen(f"task:{task_id}:stream")
+            try:
+                stream_len = await self._redis_client.xlen(f"task:{task_id}:stream")
+            except RedisError:
+                async for out in self._fatal_close(
+                    task_id, StreamErrorCode.REDIS_UNAVAILABLE.value, "redis unavailable"
+                ):
+                    yield out
+                return
             if stream_len > 0:
                 async for resp in self._consume_guarded(task_id, from_seq):
                     yield resp
@@ -436,10 +432,17 @@ class GatewayServicer:
         # First message's data is the query.
         input_key = f"task:{task_id}:input"
         if first_msg.data and len(first_msg.data.fields) > 0:
-            await self._redis_client.xadd(
-                input_key,
-                {"pb": first_msg.data.SerializeToString()},
-            )
+            try:
+                await self._redis_client.xadd(
+                    input_key,
+                    {"pb": first_msg.data.SerializeToString()},
+                )
+            except RedisError:
+                async for out in self._fatal_close(
+                    task_id, StreamErrorCode.REDIS_UNAVAILABLE.value, "redis unavailable"
+                ):
+                    yield out
+                return
 
         upstream_task = self._spawn(
             self._read_peer_upstream(request_iterator, task_id, session),
@@ -703,37 +706,150 @@ class GatewayServicer:
                 )
                 yield self._sentinel(0, task_id, "stream.end")
                 return
+            except RedisError:
+                logger.warning(
+                    "[VALIDATE R2] stream read Redis failure — closing stream: task_id=%s",
+                    task_id,
+                )  # TODO(validate): remove after prod validation
+                yield self._sentinel(
+                    0,
+                    task_id,
+                    "stream.error",
+                    code=StreamErrorCode.REDIS_UNAVAILABLE.value,
+                    message="redis unavailable during stream read",
+                    fatal=True,
+                )
+                yield self._sentinel(0, task_id, "stream.end")
+                return
             yield resp
 
-    async def _dial_consumer(  # noqa: C901, PLR0912, PLR0914, PLR0915
+    async def _dial_consumer(  # noqa: C901
         self,
         task_id: str,
         mission_id: str,
         setup_id: str,
         address: str,
-        *,
-        resume: bool = False,
     ) -> None:
-        """Dial the consumer's GatewayService.Stream and run the BiDi.
-
-        Fresh dial: sends ``stream.init``, then on the consumer's first reply
-        (the query) invokes the ``ModuleRunner`` and drains module outputs.
-        Resume dial: evicts any stale channel, sends ``stream.resume``, reads
-        the consumer's cursor from the first reply's ``from_seq``, skips the
-        runner (module already running/done), and drains from that cursor.
-        Channel is pooled by ``GrpcClientWrapper``.
+        """Dial the consumer's GatewayService.Stream, with server-side auto-reconnect.
+
+        First attempt is a fresh dial (``stream.init`` → consumer query →
+        ``ModuleRunner`` → drain from seq 0). If the BiDi dies before the stream
+        is fully delivered and the module is still producing, re-dial the SAME
+        address in resume mode (``stream.resume`` → consumer replies with its
+        ``from_seq`` → drain from that cursor, deduping the Redis queue) with
+        jittered backoff until the client returns or the reconnect window
+        (``DIGITALKIN_GATEWAY_DIAL_BACK_RECONNECT_WINDOW_S``) elapses. The module
+        runs as a separate task, spawned exactly once, so it keeps writing to
+        Redis across reconnects.
 
         Args:
             task_id: Task to push.
             mission_id: Mission ID (logging context).
             setup_id: Setup ID (logging context).
             address: ``host:port`` of the consumer's GatewayService.
-            resume: Re-attach to an existing task's output instead of starting it.
         """
         log_extra = {"task_id": task_id, "setup_id": setup_id, "mission_id": mission_id}
         # Bind IDs so the dial-back BiDi + logs carry them (task-local spawned context).
         RequestContext.bind(task_id=task_id, setup_id=setup_id, mission_id=mission_id)
 
+        module_spawned = False
+
+        def _mark_spawned() -> None:
+            nonlocal module_spawned
+            module_spawned = True
+
+        reconnect = get_gateway_settings().dial_reconnect
+        deadline: float | None = None
+        attempt = 0
+        resume = False
+        try:
+            while True:
+                disconnected = await self._run_dial_attempt(
+                    task_id=task_id,
+                    mission_id=mission_id,
+                    setup_id=setup_id,
+                    address=address,
+                    resume=resume,
+                    on_runner_spawn=_mark_spawned,
+                )
+                if not disconnected:
+                    break
+                # Client BiDi died mid-stream. Re-dial only while there is a
+                # retained stream to resume and the reboot window hasn't elapsed.
+                if not module_spawned:
+                    break
+                session = self._registry.get(task_id)
+                if session is None or session._stop_event.is_set():  # noqa: SLF001
+                    break
+                try:
+                    stream_len = await self._redis_client.xlen(f"task:{task_id}:stream")
+                except RedisError:
+                    break
+                if stream_len == 0:
+                    break
+                now = time.monotonic()
+                if deadline is None:
+                    deadline = now + reconnect.window_s
+                if now >= deadline:
+                    logger.info(
+                        "[VALIDATE RD1] dial-back gave up after %.0fs reconnect window: task_id=%s attempts=%d",
+                        reconnect.window_s,
+                        task_id,
+                        attempt,
+                    )  # TODO(validate): remove after prod validation
+                    break
+                attempt += 1
+                delay = min(random.uniform(reconnect.backoff_base_s, reconnect.backoff_max_s), deadline - now)  # noqa: S311
+                logger.info(
+                    "[VALIDATE RD1] dial-back disconnected — re-dial attempt %d in %.1fs (resume): task_id=%s",
+                    attempt,
+                    delay,
+                    task_id,
+                )  # TODO(validate): remove after prod validation
+                await asyncio.sleep(max(0.0, delay))
+                resume = True
+        finally:
+            # End-of-stream cleanup (once). Output stream is left intact for replay.
+            try:
+                removed = await self._registry.unregister(task_id)
+                if removed is not None:
+                    await removed.teardown()
+            except Exception:
+                logger.exception("end-of-stream unregister failed", extra=log_extra)
+
+    async def _run_dial_attempt(  # noqa: C901, PLR0912, PLR0914, PLR0915
+        self,
+        *,
+        task_id: str,
+        mission_id: str,
+        setup_id: str,
+        address: str,
+        resume: bool,
+        on_runner_spawn: Callable[[], None],
+    ) -> bool:
+        """Run one dial-back BiDi attempt (fresh or resume).
+
+        Fresh (``resume=False``): sends ``stream.init``, spawns the
+        ``ModuleRunner`` on the consumer's first reply (calling
+        ``on_runner_spawn``), drains from seq 0. Resume: sends ``stream.resume``,
+        reads the consumer's cursor from the first reply's ``from_seq``, skips
+        the runner, drains from that cursor. Releases the channel before
+        returning; does NOT unregister the session (the caller owns lifecycle).
+
+        Args:
+            task_id: Task to push.
+            mission_id: Mission ID (logging context).
+            setup_id: Setup ID (logging context).
+            address: ``host:port`` of the consumer's GatewayService.
+            resume: Re-attach to an existing task's output instead of starting it.
+            on_runner_spawn: Called once, when this attempt spawns the module runner.
+
+        Returns:
+            True if the BiDi died before delivering ``stream.end`` (a re-dial
+            candidate); False on clean completion or a terminal failure.
+        """
+        log_extra = {"task_id": task_id, "setup_id": setup_id, "mission_id": mission_id}
+
         async def _fail(code: str, message: str) -> None:
             """Emit a fatal to Redis (fresh dial) or log only (resume).
 
@@ -758,7 +874,7 @@ class GatewayServicer:
                 StreamErrorCode.DIAL_BACK_INTERNAL.value,
                 "gateway has no client_config to dial back with",
             )
-            return
+            return False
         # setup_version_id is unused on the dial-back path.
         comm = GrpcCommunication(
             mission_id=mission_id,
@@ -780,14 +896,14 @@ class GatewayServicer:
                 StreamErrorCode.DIAL_BACK_UNREACHABLE.value,
                 f"dial-back channel build failed: {exc}",
             )
-            return
+            return False
         except OSError as exc:
             logger.exception("dial_consumer: channel build failed addr=%s", address, extra=log_extra)
             await _fail(
                 StreamErrorCode.DIAL_BACK_UNREACHABLE.value,
                 f"dial-back channel build failed: {type(exc).__name__}: {exc}",
             )
-            return
+            return False
         t_stub = time.perf_counter_ns()
         logger.info("→ Dial-back channel ready to %s", address, extra=log_extra)
 
@@ -820,7 +936,7 @@ class GatewayServicer:
                 extra=log_extra,
             )
             await release()
-            return
+            return False
 
         # Outbound is StreamServer, inbound is StreamClient — both share
         # field tags so re-wrapping ``_consume_from_redis`` output is a rename.
@@ -835,8 +951,11 @@ class GatewayServicer:
         outgoing_done = asyncio.Event()
         # Consumer's resume cursor, captured from the first reply's ``from_seq``.
         resume_cursor = 0
+        # Set once the reader reaches EOS and stream.end is delivered to the consumer.
+        delivered_eos = False
 
         async def _outgoing() -> AsyncGenerator:
+            nonlocal delivered_eos
             try:
                 yield init_server
                 logger.info(
@@ -857,6 +976,7 @@ class GatewayServicer:
                     try:
                         cli_msg = await asyncio.wait_for(anext(reader_iter), timeout=idle_timeout)
                     except StopAsyncIteration:
+                        delivered_eos = True
                         break
                     except asyncio.TimeoutError:
                         logger.warning(
@@ -869,6 +989,25 @@ class GatewayServicer:
                             f"dial-back idle timeout: no module output in {idle_timeout:.0f}s",
                         )
                         return
+                    except RedisError:
+                        logger.warning(
+                            "[VALIDATE R2] dial-back stream read Redis failure — closing BiDi: task_id=%s",
+                            task_id,
+                            extra=log_extra,
+                        )  # TODO(validate): remove after prod validation
+                        for sc in (
+                            self._sentinel(
+                                0,
+                                task_id,
+                                "stream.error",
+                                code=StreamErrorCode.REDIS_UNAVAILABLE.value,
+                                message="redis unavailable during stream read",
+                                fatal=True,
+                            ),
+                            self._sentinel(0, task_id, "stream.end"),
+                        ):
+                            yield gateway_pb2.StreamServer(task_id=task_id, seq=sc.from_seq, data=sc.data)
+                        return
                     yield gateway_pb2.StreamServer(
                         task_id=task_id,
                         seq=cli_msg.from_seq,
@@ -880,6 +1019,7 @@ class GatewayServicer:
         async def _runner_fatal(code: str, message: str) -> None:
             await self._emit_fatal_to_redis(task_id, code=code, message=message, log_extra=log_extra)
 
+        retriable = False
         t_pre_stream = time.perf_counter_ns()
         logger.info(
             "[dial-debug] pre_stream dt_since_ready=%.3fms ch_state=%s",
@@ -953,7 +1093,7 @@ class GatewayServicer:
                             "gateway has no ModuleRunner configured",
                         )
                         output_started.set()
-                        return
+                        return False
                     self._spawn(
                         self._module_runner.run(
                             upstream.data,
@@ -964,14 +1104,23 @@ class GatewayServicer:
                         ),
                         name=f"module_runner_{task_id}",
                     )
+                    on_runner_spawn()
                     output_started.set()
                     first = False
                     continue
                 # Follow-up multi-turn input → task's input stream.
-                await self._redis_client.xadd(
-                    f"task:{task_id}:input",
-                    {"pb": upstream.data.SerializeToString()},
-                )
+                try:
+                    await self._redis_client.xadd(
+                        f"task:{task_id}:input",
+                        {"pb": upstream.data.SerializeToString()},
+                    )
+                except RedisError:
+                    logger.warning(
+                        "[VALIDATE R2] dropped follow-up input (redis unavailable): task_id=%s data=%s",
+                        task_id,
+                        str(upstream.data)[:512],
+                        extra=log_extra,
+                    )  # TODO(validate): remove after prod validation
         except grpc.aio.AioRpcError as exc:
             code_name = exc.code().name
             details = exc.details() or ""
@@ -997,6 +1146,8 @@ class GatewayServicer:
                 )
                 # Suppress DIAL_BACK_NO_QUERY in finally — RPC error already emitted.
                 output_started.set()
+            # Client BiDi died before the stream finished → re-dial candidate.
+            retriable = not delivered_eos
         except _GrpcUsageError:
             t_fail = time.perf_counter_ns()
             logger.warning(
@@ -1013,6 +1164,7 @@ class GatewayServicer:
                     "dial-back channel closed before BiDi could start",
                 )
                 output_started.set()
+            retriable = not delivered_eos
         except (RuntimeError, AssertionError, ValueError):
             logger.exception("dial_consumer unexpected error", extra=log_extra)
             if not output_started.is_set():
@@ -1035,10 +1187,4 @@ class GatewayServicer:
             # Unblock _outgoing if consumer never replied.
             output_started.set()
             await release()
-            # End-of-stream cleanup; output stream is left intact for replay.
-            try:
-                removed = await self._registry.unregister(task_id)
-                if removed is not None:
-                    await removed.teardown()
-            except Exception:
-                logger.exception("end-of-stream unregister failed", extra=log_extra)
+        return retriable

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev20}/src/digitalkin/grpc_servers/utils/circuit_breaker.py

@@ -134,6 +134,22 @@ class CircuitBreaker:
                 self._failure_count,
             )
 
+    def release_probe(self) -> bool:
+        """Release a half-open probe slot without recording an outcome.
+
+        Called when an in-flight probe is abandoned (e.g. the caller was
+        cancelled) so the breaker cannot wedge with the lock held. A cancelled
+        probe is neither success nor failure — the slot simply frees for the
+        next caller. No-op unless a probe lock is currently held.
+
+        Returns:
+            True if a held probe lock was released; False if none was held.
+        """
+        if self._state == CBState.HALF_OPEN and self._half_open_lock:
+            self._half_open_lock = False
+            return True
+        return False
+
     def reset(self) -> None:
         """Force reset to CLOSED state."""
         self._state = CBState.CLOSED