digitalkin 1.0.0.dev8__tar.gz → 1.0.0.dev9__tar.gz

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 1.0.0.dev8
+Version: 1.0.0.dev9
 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.dev8 → digitalkin-1.0.0.dev9}/pyproject.toml

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

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/__version__.py

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

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/core/task_manager/redis/redis_client.py

@@ -51,11 +51,13 @@ class RedisClient:  # noqa: PLR0904
             self.url,
             max_connections=default_size,
             decode_responses=False,
+            health_check_interval=pool.health_check_interval,
         )
         self._blocking_client = aioredis.Redis.from_url(
             self.url,
             max_connections=blocking_size,
             decode_responses=False,
+            health_check_interval=pool.health_check_interval,
         )
 
         from digitalkin.grpc_servers.utils.validators import GatewayValidator

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/core/task_manager/redis/redis_signal.py

@@ -7,6 +7,7 @@ import contextlib
 import json
 import random
 import time
+import uuid
 from typing import TYPE_CHECKING, Any, ClassVar
 
 from digitalkin.core.resilience.task_supervisor import log_unhandled
@@ -15,12 +16,20 @@ from digitalkin.logger import logger
 from digitalkin.models.settings.redis import get_redis_settings
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
     from digitalkin.core.task_manager.task_session import TaskSession
 
+    CacheInvalidator = Callable[[str, str], Awaitable[None]]
+
 
 class SharedRedisListener:
     """One PubSub connection per Redis URL; direct-dispatches signals to tasks."""
 
+    PROCESS_ID: ClassVar[str] = uuid.uuid4().hex
+    """Per-process UUID generated at class definition; identifies this listener on
+    ``signal_ch:_global_`` broadcasts. ``os.getpid()`` collides in Docker (always 1)."""
+
     _instances: ClassVar[dict[str, SharedRedisListener]] = {}
 
     @classmethod
@@ -74,6 +83,7 @@ class SharedRedisListener:
         self._pubsub: Any = None
         self._listen_task: asyncio.Task[None] | None = None
         self._stop_event = asyncio.Event()
+        self._start_lock = asyncio.Lock()
         self._counters: dict[str, int] = {
             "received": 0,
             "deduped": 0,
@@ -81,58 +91,70 @@ class SharedRedisListener:
             "dropped": 0,
             "restarts": 0,
             "subscribed": 0,
+            "invalidated": 0,
         }
         self._last_counters_log = time.monotonic()
+        self._cache_invalidator: CacheInvalidator | None = None
+
+    def set_cache_invalidator(self, handler: CacheInvalidator) -> None:
+        """Register the ``(action_name, setup_id)`` handler invoked for ``invalidate_*`` signals."""
+        self._cache_invalidator = handler
+
+    async def start(self) -> None:
+        """Open PubSub, PSUBSCRIBE ``signal_ch:*``, and start the listen loop. Idempotent under concurrent callers."""
+        async with self._start_lock:
+            if self._pubsub is not None and self._listen_task is not None and not self._listen_task.done():
+                return
+            psub_t0 = time.perf_counter_ns()
+            self._pubsub = self._redis_client.pubsub()
+            await self._pubsub.psubscribe("signal_ch:*")
+            psub_ms = (time.perf_counter_ns() - psub_t0) / 1e6
+            logger.info(
+                "[lat-audit] signal_psubscribe: psubscribe_ms=%.2f pattern=signal_ch:* phase=boot origin=%s",
+                psub_ms, SharedRedisListener.PROCESS_ID,
+            )
+            self._stop_event = asyncio.Event()
+            self._listen_task = asyncio.create_task(self._listen_loop(), name="shared_redis_listener")
+            self._listen_task.add_done_callback(log_unhandled)
 
-    async def register(
+    def register(
         self,
         task_id: str,
         session: TaskSession,
         task: asyncio.Task[None],
     ) -> None:
-        """Subscribe ``signal_ch:{task_id}``, store session/task refs, start the listen loop.
+        """Store session + task refs; sub-millisecond, never awaits.
 
         Raises:
-            RuntimeError: If max registered tasks is exceeded.
+            RuntimeError: If max registered tasks is exceeded or ``start()`` was never called.
         """
+        if self._listen_task is None or self._listen_task.done():
+            msg = "SharedRedisListener.register called before start()"
+            raise RuntimeError(msg)
         sig = get_redis_settings().signal
         if len(self._task_refs) >= sig.max_tasks:
             msg = f"SharedRedisListener: max tasks ({sig.max_tasks}) exceeded"
             raise RuntimeError(msg)
 
-        sub_t0 = time.perf_counter_ns()
-        if self._pubsub is None:
-            self._pubsub = self._redis_client.pubsub()
-        await self._pubsub.subscribe(f"signal_ch:{task_id}")
-        sub_ms = (time.perf_counter_ns() - sub_t0) / 1e6
-
+        reg_t0 = time.perf_counter_ns()
         self._task_sessions[task_id] = session
         self._task_refs[task_id] = task
         task.add_done_callback(lambda _: self.unregister(task_id))
 
-        started_loop = False
-        if self._listen_task is None or self._listen_task.done():
-            self._stop_event = asyncio.Event()
-            self._listen_task = asyncio.create_task(self._listen_loop(), name="shared_redis_listener")
-            self._listen_task.add_done_callback(log_unhandled)
-            started_loop = True
-
         self._counters["subscribed"] += 1
         logger.info(
-            "[lat-audit] signal_subscribe: subscribe_ms=%.2f listen_loop_started=%s active_subs=%d task_id=%s",
-            sub_ms,
-            started_loop,
+            "[lat-audit] signal_subscribe: register_ms=%.2f active_subs=%d task_id=%s origin=%s",
+            (time.perf_counter_ns() - reg_t0) / 1e6,
             len(self._task_refs),
             task_id,
+            SharedRedisListener.PROCESS_ID,
         )
 
     def unregister(self, task_id: str) -> None:
-        """Drop the task_id; stop the listen loop when empty."""
+        """Drop the task_id. Loop lifetime is process-wide; ``close()`` is the only stop site."""
         self._task_refs.pop(task_id, None)
         self._task_sessions.pop(task_id, None)
         self._last_seen.pop(task_id, None)
-        if not self._task_refs:
-            self._stop_event.set()
 
     def dispatch_signal(self, task_id: str, data: dict[str, Any], raw_json: str) -> bool:
         """Route a signal: ``cancel``/``stop`` → side channel + ``task.cancel()``; other actions → audit-only.
@@ -151,6 +173,24 @@ class SharedRedisListener:
         e2e_ms = (time.time_ns() - pub_ns) / 1e6 if pub_ns else 0.0
         self._counters["received"] += 1
 
+        if action.startswith("invalidate_"):
+            origin = data.get("origin")
+            if origin is not None and origin == SharedRedisListener.PROCESS_ID:
+                return True
+            setup_id = data.get("setup_id", "")
+            self._counters["invalidated"] += 1
+            logger.info(
+                "[lat-audit] signal_invalidate: e2e_ms=%.2f action=%s setup_id=%s",
+                e2e_ms, action, setup_id,
+            )
+            if self._cache_invalidator is not None:
+                inv_task = asyncio.create_task(
+                    self._cache_invalidator(action.upper(), setup_id),
+                    name=f"invalidate_{action}",
+                )
+                inv_task.add_done_callback(log_unhandled)
+            return True
+
         logger.info(
             "[lat-audit] signal_dispatch: e2e_ms=%.2f dispatch_ms=%.2f action=%s task_id=%s",
             e2e_ms,
@@ -184,7 +224,7 @@ class SharedRedisListener:
         Returns:
             The triple, or ``None`` if the message is not a usable ``signal_ch:`` payload.
         """
-        if msg["type"] != "message":
+        if msg["type"] not in {"message", "pmessage"}:
             return None
         channel = msg["channel"].decode() if isinstance(msg["channel"], bytes) else msg["channel"]
         if not channel.startswith("signal_ch:"):
@@ -204,8 +244,14 @@ class SharedRedisListener:
         while not self._stop_event.is_set():
             try:
                 if self._pubsub is None:
-                    await asyncio.sleep(0.05)
-                    continue
+                    self._pubsub = self._redis_client.pubsub()
+                    psub_t0 = time.perf_counter_ns()
+                    await self._pubsub.psubscribe("signal_ch:*")
+                    psub_ms = (time.perf_counter_ns() - psub_t0) / 1e6
+                    logger.info(
+                        "[lat-audit] signal_psubscribe: psubscribe_ms=%.2f pattern=signal_ch:* phase=loop",
+                        psub_ms,
+                    )
                 msg = await self._pubsub.get_message(ignore_subscribe_messages=True, timeout=0.5)
                 if msg is not None:
                     parsed = self._parse_message(msg)
@@ -225,6 +271,7 @@ class SharedRedisListener:
                 break
             except Exception:
                 self._counters["restarts"] += 1
+                self._pubsub = None
                 logger.exception("SharedRedisListener iteration error, retrying in %.1fs", backoff)
                 await asyncio.sleep(backoff)
                 backoff = min(backoff * 2, 10.0)
@@ -234,7 +281,8 @@ class SharedRedisListener:
                 c = self._counters
                 logger.info(
                     "[lat-audit] signal_counters: received=%d deduped=%d evicted=%d "
-                    "dropped=%d listener_restarts=%d active_subs=%d subscribed_total=%d",
+                    "dropped=%d listener_restarts=%d active_subs=%d subscribed_total=%d "
+                    "invalidated=%d",
                     c["received"],
                     c["deduped"],
                     c["evicted"],
@@ -242,6 +290,7 @@ class SharedRedisListener:
                     c["restarts"],
                     len(self._task_refs),
                     c["subscribed"],
+                    c["invalidated"],
                 )
                 self._last_counters_log = now
         self._listen_task = None
@@ -258,7 +307,7 @@ class SharedRedisListener:
         self._last_seen.clear()
         if self._pubsub is not None:
             with contextlib.suppress(Exception):
-                await self._pubsub.unsubscribe()
+                await self._pubsub.punsubscribe("signal_ch:*")
             with contextlib.suppress(Exception):
                 await self._pubsub.aclose()
             self._pubsub = None

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/core/task_manager/task_executor.py

@@ -126,7 +126,7 @@ class TaskExecutor:
                 )
             else:
                 try:
-                    await listener.register(task_id, session, task)
+                    listener.register(task_id, session, task)
                 except Exception:
                     logger.warning(
                         "Signal registration failed — signals disabled for task_id=%s",

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/grpc_servers/gateway_servicer.py

@@ -17,6 +17,7 @@ from redis.exceptions import RedisError
 
 from digitalkin.core.profiling.step_timer import StepTimer
 from digitalkin.core.task_manager.redis.proto_streams import ProtoStreamReader
+from digitalkin.core.task_manager.redis.redis_signal import SharedRedisListener
 from digitalkin.grpc_servers.m2m_call_registry import M2MCallRegistry
 from digitalkin.grpc_servers.stream_registry import StreamRegistry
 from digitalkin.grpc_servers.stream_session import StreamSession
@@ -130,8 +131,17 @@ class GatewayServicer:
         return task
 
     async def start(self) -> None:
-        """Start the M2M call-registry TTL sweeper."""
+        """Start the M2M call-registry TTL sweeper and PSUBSCRIBE the signal listener."""
         await self._m2m.start()
+        listener = SharedRedisListener.singleton_or_none()
+        if listener is not None:
+            try:
+                await listener.start()
+            except Exception:
+                logger.warning(
+                    "SharedRedisListener.start() failed at boot — first-task PSUBSCRIBE will retry lazily",
+                    exc_info=True,
+                )
 
     async def stop(self) -> None:
         """Shut down registries and cancel the M2M sweeper."""
@@ -433,26 +443,35 @@ class GatewayServicer:
 
         try:
             if action_name.startswith("INVALIDATE_"):
+                setup_id_for_invalidate = task_id
                 if self._cache_handler is not None:
-                    await self._cache_handler(action_name)
+                    await self._cache_handler(action_name, setup_id_for_invalidate)
                     timer.mark("cache_handler")
                     last_mark = "cache_handler"
-                    logger.info(
-                        "[lat-audit] SendSignal: %s path=cache total=%.2fms action=%s",
-                        timer.format_steps(),
-                        timer.total_ms(),
-                        action_name,
+                payload = json.dumps({
+                    "action": action_name.lower(),
+                    "setup_id": setup_id_for_invalidate,
+                    "published_at_ns": time.time_ns(),
+                })
+                try:
+                    await self._redis_client.publish("signal_ch:_global_", payload)
+                    timer.mark("global_publish")
+                    last_mark = "global_publish"
+                except RedisError:
+                    logger.warning(
+                        "[gateway] INVALIDATE fan-out publish failed — local-only invalidation applied",
                         extra=log_extra,
+                        exc_info=True,
                     )
-                    return gateway_pb2.ClientSignalResponse(success=True, task_id="")
-                logger.warning(
-                    "[gateway] SendSignal_failed: failure=NoCacheHandler at_step=%s elapsed_ms=%.2f action=%s",
-                    last_mark,
-                    timer.elapsed_now_ms(),
+                logger.info(
+                    "[lat-audit] SendSignal: %s path=cache total=%.2fms action=%s setup_id=%s",
+                    timer.format_steps(),
+                    timer.total_ms(),
                     action_name,
+                    setup_id_for_invalidate,
                     extra=log_extra,
                 )
-                return gateway_pb2.ClientSignalResponse(success=False, task_id="")
+                return gateway_pb2.ClientSignalResponse(success=True, task_id=setup_id_for_invalidate)
 
             if GatewayValidator.validate_id(task_id, "task_id") is not None:
                 logger.warning(

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/grpc_servers/module_server.py

@@ -9,6 +9,7 @@ from agentic_mesh_protocol.module.v1 import module_service_pb2, module_service_p
 
 from digitalkin.core.task_manager.module_runner import ModuleRunner
 from digitalkin.core.task_manager.redis import RedisClient
+from digitalkin.core.task_manager.redis.redis_signal import SharedRedisListener
 from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.gateway_servicer import GatewayServicer
 from digitalkin.grpc_servers.module_servicer import ModuleServicer
@@ -129,11 +130,20 @@ class ModuleServer(BaseServer):
 
         logger.info("GatewayServicer + ModuleRunner registered (Redis: %s)", redis_url)
 
-    async def _handle_cache_invalidation(self, action: str) -> None:
-        """Dispatch cache invalidation by action name. Isolated from module lifecycle.
+        listener = SharedRedisListener.singleton_or_none()
+        if listener is not None:
+            listener.set_cache_invalidator(self._handle_cache_invalidation)
+
+    async def _handle_cache_invalidation(self, action: str, setup_id: str = "") -> None:
+        """Dispatch cache invalidation by action name.
+
+        ``INVALIDATE_SETUP`` and ``INVALIDATE_TOOLS`` require a ``setup_id`` —
+        without one they log a warning and skip. ``INVALIDATE_ALL`` is the only
+        full-wipe path.
 
         Args:
-            action: SignalAction enum name (e.g. "INVALIDATE_SETUP").
+            action: SignalAction enum name (e.g. ``INVALIDATE_SETUP``).
+            setup_id: Setup identifier for scoped invalidation; ignored by full-wipe actions.
         """
         handlers: dict[str, Any] = {
             "INVALIDATE_ALL": self._invalidate_all,
@@ -144,24 +154,39 @@ class ModuleServer(BaseServer):
             "INVALIDATE_SHARED": self._invalidate_shared,
         }
         handler = handlers.get(action)
-        if handler is not None:
+        if handler is None:
+            logger.warning("Unknown invalidation action: %s", action)
+            return
+        if action in {"INVALIDATE_SETUP", "INVALIDATE_TOOLS"}:
+            await handler(setup_id)
+        else:
             await handler()
-            logger.info("Cache invalidated: %s", action)
+        logger.info("Cache invalidated: %s setup_id=%s", action, setup_id or "<all>")
 
     async def _invalidate_all(self) -> None:
-        await self._invalidate_setup()
-        await self._invalidate_tools()
+        if self.module_servicer is not None:
+            self.module_servicer.invalidate_setup_cache()
+            self.module_servicer.invalidate_tool_cache()
         await self._invalidate_shared()
         await self._invalidate_models()
         await self._invalidate_channels()
 
-    async def _invalidate_setup(self) -> None:
-        if self.module_servicer is not None:
-            self.module_servicer.invalidate_setup_cache()
+    async def _invalidate_setup(self, setup_id: str = "") -> None:
+        if self.module_servicer is None:
+            return
+        if not setup_id:
+            logger.warning("INVALIDATE_SETUP received without setup_id — skipping (scoped-only policy)")
+            return
+        self.module_servicer._setup_cache.pop(setup_id, None)  # noqa: SLF001
+        self.module_servicer._setup_inflight.pop(setup_id, None)  # noqa: SLF001
 
-    async def _invalidate_tools(self) -> None:
-        if self.module_servicer is not None:
-            self.module_servicer.invalidate_tool_cache()
+    async def _invalidate_tools(self, setup_id: str = "") -> None:
+        if self.module_servicer is None:
+            return
+        if not setup_id:
+            logger.warning("INVALIDATE_TOOLS received without setup_id — skipping (scoped-only policy)")
+            return
+        self.module_servicer._tool_cache_by_setup.pop(setup_id, None)  # noqa: SLF001
 
     async def _invalidate_shared(self) -> None:
         self.module_class.clear_shared()

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/grpc_servers/module_servicer.py

@@ -1,6 +1,7 @@
 """Module servicer implementation for DigitalKin."""
 
 import asyncio
+import json
 import os
 import time
 from argparse import ArgumentParser, Namespace
@@ -17,6 +18,7 @@ from google.protobuf import json_format, struct_pb2
 
 from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.core.job_manager.single_job_manager import SingleJobManager
+from digitalkin.core.task_manager.redis.redis_signal import SharedRedisListener
 from digitalkin.grpc_servers.exceptions import ServicerError
 from digitalkin.logger import logger
 from digitalkin.models.module.module import ModuleCodeModel
@@ -387,6 +389,26 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             ),
         )
         self._tool_cache_by_setup.pop(setup_version.setup_id, None)
+
+        publish_ns = time.time_ns()
+        for action in ("invalidate_setup", "invalidate_tools"):
+            payload = json.dumps({
+                "action": action,
+                "setup_id": setup_version.setup_id,
+                "published_at_ns": publish_ns,
+                "origin": SharedRedisListener.PROCESS_ID,
+            })
+            try:
+                await self._redis_client.publish("signal_ch:_global_", payload)
+            except Exception:
+                logger.warning(
+                    "[gateway] cache-invalidate fan-out publish failed for action=%s "
+                    "setup_id=%s — peers may keep stale cache until TTL",
+                    action,
+                    setup_version.setup_id,
+                    exc_info=True,
+                )
+
         setup_version.content = json_format.ParseDict(  # type: ignore[misc]
             updated_setup_data,
             struct_pb2.Struct(),

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin/models/settings/redis.py

@@ -16,6 +16,10 @@ class RedisPoolSettings(BaseSettings):
     pool_size_default: int = Field(default=0, description="Non-blocking pool size (0 = pool_size // 2)")
     pool_size_blocking: int = Field(default=0, description="Blocking pool size for XREAD (0 = pool_size // 2)")
     health_check_timeout: float = Field(default=5.0, description="Max seconds to wait for a PING during health check.")
+    health_check_interval: int = Field(
+        default=15,
+        description="Seconds between connection-level PINGs; 0 disables. Catches silently-dead sockets.",
+    )
 
     def get_default_pool_size(self) -> int:
         """Non-blocking pool size, defaults to half of total.

{digitalkin-1.0.0.dev8 → digitalkin-1.0.0.dev9}/src/digitalkin.egg-info/PKG-INFO

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