digitalkin 1.0.0.dev18__tar.gz → 1.0.0.dev19__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 1.0.0.dev18
+Version: 1.0.0.dev19
 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.dev19}/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.dev19"
 
     [project.optional-dependencies]
         performance = [ "uvloop>=0.21" ]

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev19}/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.dev19"

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev19}/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.dev19}/src/digitalkin/grpc_servers/gateway_servicer.py

@@ -157,7 +157,7 @@ class GatewayServicer:
         await self._m2m.stop()
         await self._registry.shutdown()
 
-    async def StartStream(  # noqa: PLR0911, PLR0915
+    async def StartStream(  # noqa: C901, PLR0911, PLR0915
         self,
         request: Any,
         context: grpc.aio.ServicerContext,
@@ -215,14 +215,30 @@ class GatewayServicer:
         # 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)
+        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")
+            try:
+                stream_len = await self._redis_client.xlen(f"task:{task_id}:stream")
+            except RedisError:
+                logger.warning(
+                    "[VALIDATE R2] StartStream xlen 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)
             if stream_len == 0:
                 logger.info(
                     "[VALIDATE IDEM1] no live stream to resume (TTL expired) task_id=%s (%s)",
@@ -274,7 +290,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 +306,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 +453,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 +472,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,6 +746,21 @@ 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
@@ -869,6 +927,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,
@@ -968,10 +1045,18 @@ class GatewayServicer:
                     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 ""

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev19}/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

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev19}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py

@@ -224,7 +224,7 @@ class GrpcClientWrapper:
         cls._stub_cache.clear()
         CircuitBreaker.clear_all()
 
-    async def exec_grpc_query(  # noqa: C901, PLR0914
+    async def exec_grpc_query(  # noqa: C901, PLR0912, PLR0914, PLR0915
         self,
         query_endpoint: str,
         request: Any,
@@ -284,57 +284,66 @@ class GrpcClientWrapper:
         backoff_delays = tuple(backoff_base_ms / 1000 * (2**i) for i in range(max_retries))
         last_error: grpc.RpcError | None = None
 
-        for attempt in range(max_retries + 1):
-            if attempt > 0:
-                await asyncio.sleep(backoff_delays[attempt - 1])
-
-            try:
-                response = await rpc_method(request, timeout=eff_timeout)
-            except grpc.RpcError as e:
-                last_error = e
-                if e.code() in self._RETRYABLE_CODES and attempt < max_retries:
-                    logger.warning(
-                        "gRPC transient error on %s.%s [%s] (attempt %d/%d), retrying in %.0fms",
-                        self.service_name,
-                        query_endpoint,
-                        e.code().name,
-                        attempt + 1,
-                        max_retries + 1,
-                        backoff_delays[attempt] * 1000,
-                    )
-                    continue
-                if e.code() in self._CIRCUIT_FAILURE_CODES:
-                    logger.warning(
-                        "circuit-breaker tick: %s.%s [%s]",
-                        self.service_name,
-                        query_endpoint,
-                        e.code().name,
-                    )
-                    cb.record_failure()
+        try:
+            for attempt in range(max_retries + 1):
+                if attempt > 0:
+                    await asyncio.sleep(backoff_delays[attempt - 1])
+
+                try:
+                    response = await rpc_method(request, timeout=eff_timeout)
+                except grpc.RpcError as e:
+                    last_error = e
+                    if e.code() in self._RETRYABLE_CODES and attempt < max_retries:
+                        logger.warning(
+                            "gRPC transient error on %s.%s [%s] (attempt %d/%d), retrying in %.0fms",
+                            self.service_name,
+                            query_endpoint,
+                            e.code().name,
+                            attempt + 1,
+                            max_retries + 1,
+                            backoff_delays[attempt] * 1000,
+                        )
+                        continue
+                    if e.code() in self._CIRCUIT_FAILURE_CODES:
+                        logger.warning(
+                            "circuit-breaker tick: %s.%s [%s]",
+                            self.service_name,
+                            query_endpoint,
+                            e.code().name,
+                        )
+                        cb.record_failure()
+                    else:
+                        cb.record_success()
+                    break
                 else:
                     cb.record_success()
-                break
-            else:
-                cb.record_success()
-                return response
-
-        if last_error is None:
-            msg = f"[gRPC-client:{self.service_name}.{query_endpoint}] Retry loop exited without response or error"
-            raise ServerError(msg)
-        status_code = last_error.code().name
-        details = last_error.details()
-        retried = last_error.code() in self._RETRYABLE_CODES
-        suffix = f" (after {max_retries + 1} attempts)" if retried else ""
-
-        log_level = logging.DEBUG if last_error.code() == grpc.StatusCode.NOT_FOUND else logging.ERROR
-        logger.log(
-            log_level,
-            "gRPC call failed: %s.%s [%s] %s (%s)",
-            self.service_name,
-            query_endpoint,
-            status_code,
-            details,
-            type(request).__name__,
-        )
-        error_msg = f"[gRPC-client:{self.service_name}.{query_endpoint}] [{status_code}] {details}{suffix}"
-        raise ServerError(error_msg) from last_error
+                    return response
+
+            if last_error is None:
+                msg = f"[gRPC-client:{self.service_name}.{query_endpoint}] Retry loop exited without response or error"
+                raise ServerError(msg)
+            status_code = last_error.code().name
+            details = last_error.details()
+            retried = last_error.code() in self._RETRYABLE_CODES
+            suffix = f" (after {max_retries + 1} attempts)" if retried else ""
+
+            log_level = logging.DEBUG if last_error.code() == grpc.StatusCode.NOT_FOUND else logging.ERROR
+            logger.log(
+                log_level,
+                "gRPC call failed: %s.%s [%s] %s (%s)",
+                self.service_name,
+                query_endpoint,
+                status_code,
+                details,
+                type(request).__name__,
+            )
+            error_msg = f"[gRPC-client:{self.service_name}.{query_endpoint}] [{status_code}] {details}{suffix}"
+            raise ServerError(error_msg) from last_error
+        finally:
+            # Free a half-open probe the loop never resolved (e.g. CancelledError mid-call or backoff).
+            if cb.release_probe():
+                logger.info(
+                    "[VALIDATE R1] released abandoned circuit probe: %s.%s",
+                    self.service_name,
+                    query_endpoint,
+                )  # TODO(validate): remove after prod validation

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev19}/src/digitalkin/models/grpc_servers/stream_error_codes.py

@@ -19,6 +19,7 @@ class StreamErrorCode(str, Enum):
     DIAL_BACK_NO_QUERY = "DIAL_BACK_NO_QUERY"
     DIAL_BACK_IDLE_TIMEOUT = "DIAL_BACK_IDLE_TIMEOUT"
     STREAM_IDLE_TIMEOUT = "STREAM_IDLE_TIMEOUT"
+    REDIS_UNAVAILABLE = "REDIS_UNAVAILABLE"
     MODULE_RUNTIME_ERROR = "MODULE_RUNTIME_ERROR"
     INPUT_VALIDATION_ERROR = "INPUT_VALIDATION_ERROR"
     BACKPRESSURE_TIMEOUT = "BACKPRESSURE_TIMEOUT"

{digitalkin-1.0.0.dev18 → digitalkin-1.0.0.dev19}/src/digitalkin.egg-info/PKG-INFO

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