digitalkin 1.0.0.dev2__tar.gz → 1.0.0.dev3__tar.gz

{digitalkin-1.0.0.dev2 → digitalkin-1.0.0.dev3}/PKG-INFO

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

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

{digitalkin-1.0.0.dev2 → digitalkin-1.0.0.dev3}/src/digitalkin/__version__.py

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

{digitalkin-1.0.0.dev2 → digitalkin-1.0.0.dev3}/src/digitalkin/core/job_manager/base_job_manager.py

@@ -1,8 +1,7 @@
 """Background module manager."""
 
 import abc
-from collections.abc import AsyncGenerator, Callable, Coroutine
-from contextlib import AbstractAsyncContextManager
+from collections.abc import Callable, Coroutine
 from typing import Any, Generic
 
 from digitalkin.core.task_manager.base_task_manager import BaseTaskManager
@@ -79,18 +78,6 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
         """
         await self._task_manager.create_task(task_id, mission_id, module, coro, **kwargs)
 
-    async def clean_session(self, task_id: str, mission_id: str) -> bool:
-        """Clean a task's session.
-
-        Args:
-            task_id: Unique identifier for the task.
-            mission_id: Mission identifier.
-
-        Returns:
-            bool: True if the task was successfully cancelled, False otherwise.
-        """
-        return await self._task_manager.clean_session(task_id, mission_id)
-
     async def cancel_task(self, task_id: str, mission_id: str, timeout: float | None = None) -> bool:
         """Cancel a task.
 
@@ -164,48 +151,6 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
 
         return callback_wrapper
 
-    @abc.abstractmethod
-    def generate_stream_consumer(
-        self, job_id: str
-    ) -> AbstractAsyncContextManager[AsyncGenerator[dict[str, Any], None]]:
-        """Generate a stream consumer for the job's message stream.
-
-        Args:
-            job_id: The unique identifier of the job to filter messages for.
-
-        Yields:
-            dict[str, Any]: The messages from the associated module's stream.
-        """
-
-    @abc.abstractmethod
-    async def create_module_instance_job(
-        self,
-        input_data: InputModelT,
-        setup_data: SetupModelT,
-        mission_id: str,
-        setup_id: str,
-        setup_version_id: str,
-        request_metadata: dict[str, str] | None = None,
-        job_id: str | None = None,
-        tool_cache: Any = None,
-    ) -> str:
-        """Create and start a new job for the module's instance.
-
-        Args:
-            input_data: The input data required to start the job.
-            setup_data: The setup configuration for the module.
-            mission_id: The mission ID associated with the job.
-            setup_id: The setup ID.
-            setup_version_id: The setup version ID associated with the module.
-            request_metadata: gRPC request metadata (headers) to forward to the module.
-            job_id: Optional externally-provided job ID (e.g., Gateway's task_id).
-                    If None, a UUID is minted internally.
-            tool_cache: Pre-resolved ToolCache to inject (skips per-request resolution).
-
-        Returns:
-            str: The unique identifier (job ID) of the created job.
-        """
-
     @abc.abstractmethod
     async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT | ModuleCodeModel:
         """Generate a stream consumer for a module's output data.
@@ -249,38 +194,6 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
             Exception: If the module fails to start.
         """
 
-    @abc.abstractmethod
-    async def stop_module(self, job_id: str) -> bool:
-        """Stop a running module job.
-
-        Args:
-            job_id: The unique identifier of the job to stop.
-
-        Returns:
-            bool: True if the job was successfully stopped, False if it does not exist.
-        """
-
-    @abc.abstractmethod
-    async def wait_for_completion(self, job_id: str) -> None:
-        """Wait for a task to complete.
-
-        This method blocks until the specified job has reached a terminal state.
-        SingleJobManager awaits the asyncio.Task directly.
-
-        Args:
-            job_id: The unique identifier of the job to wait for.
-
-        Raises:
-            KeyError: If the job_id is not found.
-        """
-
-    @abc.abstractmethod
-    async def stop_all_modules(self) -> None:
-        """Stop all currently running module jobs.
-
-        This method ensures that all active jobs are gracefully terminated.
-        """
-
     @abc.abstractmethod
     async def list_modules(self) -> dict[str, dict[str, Any]]:
         """List all modules along with their statuses.

{digitalkin-1.0.0.dev2 → digitalkin-1.0.0.dev3}/src/digitalkin/core/job_manager/single_job_manager.py

@@ -11,7 +11,6 @@ from __future__ import annotations
 import asyncio
 import os
 import uuid
-from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any
 
 import grpc
@@ -26,7 +25,7 @@ from digitalkin.models.module.base_types import DataModel, InputModelT, OutputMo
 from digitalkin.models.module.module import ModuleCodeModel
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, AsyncIterator, Callable
+    from collections.abc import Callable
 
     from digitalkin.core.task_manager.redis.redis_client import RedisClient
     from digitalkin.core.task_manager.redis.redis_streams import RedisStreamWriter
@@ -239,104 +238,8 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
                 except asyncio.QueueFull:
                     logger.warning("Queue full, rejecting new message", extra={"job_id": job_id})
 
-    @asynccontextmanager
-    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:
-        """Generate a stream consumer for a module's output data.
-
-        This method creates an asynchronous generator that streams output data
-        from a specific module job. If the module does not exist, it generates
-        an error message.
-
-        Args:
-            job_id: The unique identifier of the job.
-
-        Yields:
-            AsyncGenerator: A stream of output data or error messages.
-        """
-        if (session := self.tasks_sessions.get(job_id, None)) is None:
-
-            async def _error_gen() -> AsyncGenerator[  # noqa: RUF029
-                dict[str, Any], None
-            ]:  # Async generator type required by caller even though body uses yield
-                """Generate an error message for a non-existent module.
-
-                Yields:
-                    AsyncGenerator: A generator yielding an error message.
-                """
-                yield {
-                    "error": {
-                        "error_message": f"Module {job_id} not found",
-                        "code": grpc.StatusCode.NOT_FOUND,
-                    }
-                }
-
-            yield _error_gen()
-            return
-
-        logger.debug("Session: %s with Module %s", job_id, session.module)
-
-        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
-            """Stream output data from the module with bounded blocking.
-
-            Uses a 1-second timeout on queue.get() to periodically re-check
-            termination flags, preventing indefinite hangs when the task crashes
-            without producing output.
-
-            Termination behavior:
-            - cancelled: abort immediately (abnormal, discard remaining)
-            - stream_closed / completed / failed: drain remaining queue items, then exit
-
-            Yields:
-                dict: Output data generated by the module.
-            """
-            while True:
-                if session.cancelled:
-                    logger.debug("Stream cancelled for job %s", job_id)
-                    break
-
-                # If no more output will be produced, drain remaining items and exit
-                if session.stream_closed or session.status in {"completed", "failed"}:
-                    while not session.queue.empty():
-                        msg = session.queue.get_nowait()
-                        try:
-                            yield msg
-                        finally:
-                            session.queue.task_done()
-                    logger.debug(
-                        "Stream drained for job %s: status=%s, stream_closed=%s",
-                        job_id,
-                        session.status,
-                        session.stream_closed,
-                    )
-                    break
-
-                try:
-                    msg = await asyncio.wait_for(session.queue.get(), timeout=0.25)
-                except asyncio.TimeoutError:
-                    continue
-
-                try:
-                    yield msg
-                finally:
-                    session.queue.task_done()
-
-                if session.cancelled:
-                    break
-
-        try:
-            yield _stream()
-        finally:
-            # Write EOS to Redis Stream if writer exists (cleanup on stream close)
-            writer = self._stream_writers.pop(job_id, None) if self._stream_writers is not None else None
-            if writer is not None:
-                try:
-                    await writer.write_eos()
-                except Exception:
-                    logger.warning("Redis stream EOS write failed", extra={"job_id": job_id})
-
-    async def create_module_instance_job(
+    async def preload_instance(
         self,
-        input_data: InputModelT,
         setup_data: SetupModelT,
         mission_id: str,
         setup_id: str,
@@ -345,28 +248,32 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         job_id: str | None = None,
         tool_cache: Any = None,
         callback: Callable | None = None,
-    ) -> str:
-        """Create and start a new module job.
+    ) -> tuple[Any, str, Callable]:
+        """Phase 3.A: Build + warm a module instance without input.
+
+        Calls the factory, wires the redis task manager + callback, and
+        runs the module's ``prepare()`` (which is idempotent — the later
+        ``start()`` call short-circuits past it).
+
+        Designed so the dial-back orchestrator can pay LiteLLM/agno init
+        costs (~440 ms) in parallel with the consumer's first reply RTT.
 
         Args:
-            input_data: The input data required to start the job.
             setup_data: The setup configuration for the module.
-            mission_id: The mission ID associated with the job.
-            setup_id: The setup ID associated with the module.
-            setup_version_id: The setup Version ID associated with the module.
-            request_metadata: gRPC request metadata (headers) to forward to the module.
-            job_id: Optional externally-provided job ID (e.g., Gateway's task_id).
-            tool_cache: Pre-resolved ToolCache to inject (skips per-request resolution).
-            callback: Direct output callback (writes to Redis). If None, uses
-                internal queue path via add_to_queue.
+            mission_id: Mission ID.
+            setup_id: Setup ID.
+            setup_version_id: Setup version ID.
+            request_metadata: gRPC request metadata (headers).
+            job_id: Optional externally-provided job ID.
+            tool_cache: Pre-resolved ToolCache.
+            callback: Direct output callback. If None, the in-memory
+                queue path is wired for the eventual ``run_preloaded``.
 
         Returns:
-            str: The unique identifier (job ID) of the created job.
-
-        Raises:
-            Exception: If the module fails to start.
+            ``(module, job_id, callback)`` — pass to ``run_preloaded``.
         """
         from digitalkin.core.profiling.step_timer import StepTimer
+        from digitalkin.services.task_manager.redis_task_manager import RedisTaskManager
 
         timer = StepTimer()
         job_id = job_id or str(uuid.uuid4())
@@ -381,9 +288,6 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         )
         timer.mark("factory_create")
 
-        # Redis-backed signal service for cross-process signal delivery
-        from digitalkin.services.task_manager.redis_task_manager import RedisTaskManager
-
         module.context.task_manager = RedisTaskManager(self._redis_client)
         timer.mark("redis_task_manager")
 
@@ -394,97 +298,49 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             callback = await self.job_specific_callback(self.add_to_queue, job_id)
             timer.mark("default_callback")
 
-        await self.create_task(
-            job_id,
-            mission_id,
-            module,
-            module.start(input_data, setup_data, callback, done_callback=None),  # type: ignore[arg-type]
-        )
-        timer.mark("create_task")
-        timer.log("create_module_instance_job", task_id=job_id)
-        logger.info("Managed task started: '%s'", job_id, extra={"task_id": job_id})
-        return job_id
-
-    async def clean_session(self, task_id: str, mission_id: str) -> bool:
-        """Clean a task's session.
+        await module.prepare(setup_data, callback)
+        timer.mark("prepare")
+        timer.log("preload_instance", task_id=job_id)
+        return module, job_id, callback
 
-        Args:
-            task_id: Unique identifier for the task.
-            mission_id: Mission identifier.
-
-        Returns:
-            bool: True if the task was successfully cleaned, False otherwise.
-        """
-        return await self._task_manager.clean_session(task_id, mission_id)
+    async def run_preloaded(
+        self,
+        module: Any,
+        job_id: str,
+        mission_id: str,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        callback: Callable,
+    ) -> str:
+        """Phase 3.A: Run a pre-prepared module instance with input.
 
-    async def stop_module(self, job_id: str) -> bool:
-        """Stop a running module job.
+        ``module`` must come from :meth:`preload_instance`. Schedules
+        the run in the task manager and returns the job_id.
 
         Args:
-            job_id: The unique identifier of the job to stop.
+            module: Pre-prepared module instance.
+            job_id: Job/task ID assigned by ``preload_instance``.
+            mission_id: Mission ID for task manager scoping.
+            input_data: The first input (the query) to feed ``run()``.
+            setup_data: The setup the instance was prepared with.
+            callback: Output callback (already attached to context).
 
         Returns:
-            bool: True if the module was successfully stopped, False if it does not exist.
-
-        Raises:
-            Exception: If an error occurs while stopping the module.
+            The ``job_id`` (echoed for caller convenience).
         """
-        logger.info("Stop module requested", extra={"job_id": job_id})
-
-        logger.debug("debug:stop_module acquiring lock job_id=%s", job_id)
-        async with self._lock:
-            session = self.tasks_sessions.get(job_id)
-
-            if not session:
-                logger.warning("Session not found", extra={"job_id": job_id})
-                return False
-            try:
-                await session.module.stop()
-                await self.cancel_task(job_id, session.mission_id)
-                logger.debug(
-                    "Module stopped successfully",
-                    extra={"job_id": job_id, "mission_id": session.mission_id},
-                )
-            except Exception:
-                logger.exception("Error stopping module", extra={"job_id": job_id})
-                raise
-            else:
-                return True
-            finally:
-                # Clean up stream writer if consumer never started
-                if self._stream_writers is not None:
-                    writer = self._stream_writers.pop(job_id, None)
-                    if writer is not None:
-                        try:
-                            await writer.write_eos()
-                        except Exception:
-                            logger.debug("EOS write failed during stop", extra={"job_id": job_id})
-
-    async def wait_for_completion(self, job_id: str) -> None:
-        """Wait for a task to complete by awaiting its asyncio.Task.
-
-        Idempotent — safe to call after the task has already been cleaned up
-        (e.g. by deferred cleanup during signal cancellation).
+        from digitalkin.core.profiling.step_timer import StepTimer
 
-        Args:
-            job_id: The unique identifier of the job to wait for.
-        """
-        task = self._task_manager.tasks.get(job_id)
-        if task is None:
-            logger.debug("Task already cleaned up, skipping wait_for_completion", extra={"job_id": job_id})
-            return
-        await task
-
-    async def stop_all_modules(self) -> None:
-        """Stop all currently running module jobs."""
-        # Snapshot job IDs while holding lock
-        async with self._lock:
-            job_ids = list(self.tasks_sessions.keys())
-
-        # Release lock before calling stop_module (which has its own lock)
-        if job_ids:
-            stop_tasks = [self.stop_module(job_id) for job_id in job_ids]
-            await asyncio.gather(*stop_tasks, return_exceptions=True)
+        timer = StepTimer()
+        await self.create_task(
+            job_id,
+            mission_id,
+            module,
+            module.start(input_data, setup_data, callback, done_callback=None),
+        )
+        timer.mark("create_task")
+        timer.log("run_preloaded", task_id=job_id)
+        logger.info("Managed task started: '%s'", job_id, extra={"task_id": job_id})
+        return job_id
 
     async def list_modules(self) -> dict[str, dict[str, Any]]:
         """List all modules along with their statuses.

{digitalkin-1.0.0.dev2 → digitalkin-1.0.0.dev3}/src/digitalkin/core/profiling/task_profiler.py

@@ -10,6 +10,35 @@ from typing import Any
 
 from digitalkin.logger import logger
 
+# Phase 7.C: rotate per-task profile files to avoid unbounded growth.
+# Default keeps the most recent N profiles by mtime; configure via
+# ``DIGITALKIN_PROFILER_KEEP_N``.
+PROFILER_KEEP_N = int(os.environ.get("DIGITALKIN_PROFILER_KEEP_N", "100"))
+
+
+def _rotate_profiles(output_dir: str, keep_n: int, suffixes: tuple[str, ...]) -> None:
+    """Trim ``output_dir`` to the most recent ``keep_n`` files by mtime.
+
+    Args:
+        output_dir: Directory containing profile files.
+        keep_n: Number of files to keep. ``<= 0`` disables rotation.
+        suffixes: File extensions to include in rotation (e.g. ``(".html",)``).
+    """
+    if keep_n <= 0:
+        return
+    try:
+        candidates = [p for p in Path(output_dir).iterdir() if p.is_file() and p.suffix in suffixes]
+    except OSError:
+        return
+    if len(candidates) <= keep_n:
+        return
+    candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True)
+    for stale in candidates[keep_n:]:
+        try:
+            stale.unlink()
+        except OSError:
+            logger.debug("Profiler rotation: could not delete %s", stale)
+
 
 class ProfilerMode(str, Enum):
     """Profiler backend selection."""
@@ -121,6 +150,7 @@ class TaskProfiler:
                 Path(path).write_text(self._profiler.output_html(), encoding="utf-8")
                 logger.info("Pyinstrument profile saved: %s", path)
                 logger.info("Pyinstrument summary:\n%s", self._profiler.output_text())
+                _rotate_profiles(self._output_dir, PROFILER_KEEP_N, (".html",))
 
         except Exception:
             logger.exception("Failed to stop/save profiler %s for task %s", self._mode.value, self._task_id)

digitalkin-1.0.0.dev3/src/digitalkin/core/resilience/task_supervisor.py

@@ -0,0 +1,37 @@
+"""Tiny helper: log unhandled exceptions on fire-and-forget asyncio tasks."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from digitalkin.logger import logger
+
+
+def log_unhandled(task: asyncio.Task[Any]) -> None:
+    """Done-callback that logs uncaught exceptions on a fire-and-forget task.
+
+    Cancellation and clean exits are silent. Anything else is logged at
+    error level with the task name and traceback — this replaces asyncio's
+    opaque ``Task exception was never retrieved`` warning with an
+    actionable log line.
+
+    Usage:
+
+        task = asyncio.create_task(coro, name="my_daemon")
+        task.add_done_callback(log_unhandled)
+
+    Args:
+        task: The done asyncio task to inspect.
+    """
+    if task.cancelled():
+        return
+    exc = task.exception()
+    if exc is not None:
+        logger.error(
+            "Background task '%s' failed with %s: %s",
+            task.get_name(),
+            type(exc).__name__,
+            exc,
+            exc_info=exc,
+        )

{digitalkin-1.0.0.dev2 → digitalkin-1.0.0.dev3}/src/digitalkin/core/task_manager/base_task_manager.py

@@ -46,7 +46,6 @@ class BaseTaskManager(ABC):
         self._active_slots = 0
         self._task_wait_timeout = float(os.environ.get("DIGITALKIN_TASK_WAIT_TIMEOUT", "30"))
         self._stream_drain_timeout = float(os.environ.get("DIGITALKIN_STREAM_DRAIN_TIMEOUT", "2.0"))
-        self._cleanup_tasks: set[asyncio.Task] = set()
 
         # Admission queue: allows tasks to wait for a slot instead of being rejected.
         # Total in-system capacity = max_concurrent + max_queued.
@@ -278,49 +277,6 @@ class BaseTaskManager(ABC):
         self.tasks_sessions[task_id] = session
         return session
 
-    def _register_auto_cleanup(self, task_id: str, mission_id: str) -> None:
-        """Register a done callback on the supervisor task for deferred cleanup.
-
-        When the supervisor finishes, waits for the stream consumer to drain
-        (up to 60s), then runs idempotent cleanup. Safe if the servicer
-        already cleaned up.
-
-        Args:
-            task_id: The ID of the task.
-            mission_id: The ID of the mission.
-        """
-        supervisor = self.tasks.get(task_id)
-        if supervisor is None:
-            return
-
-        def _on_done(_: asyncio.Task) -> None:
-            t = asyncio.ensure_future(self._deferred_cleanup(task_id, mission_id))
-            self._cleanup_tasks.add(t)
-            t.add_done_callback(self._cleanup_tasks.discard)
-
-        supervisor.add_done_callback(_on_done)
-
-    async def _deferred_cleanup(self, task_id: str, mission_id: str) -> None:
-        """Wait for stream drain then cleanup.
-
-        Args:
-            task_id: The ID of the task.
-            mission_id: The ID of the mission.
-        """
-        session = self.tasks_sessions.get(task_id)
-        if session is None:
-            return
-
-        try:
-            await asyncio.wait_for(session._stream_closed.wait(), timeout=self._stream_drain_timeout)  # noqa: SLF001
-        except asyncio.TimeoutError:
-            logger.warning(
-                "Stream drain timeout, proceeding with cleanup",
-                extra={"task_id": task_id, "mission_id": mission_id},
-            )
-
-        await self._cleanup_task(task_id, mission_id)
-
     @abstractmethod
     async def create_task(
         self,
@@ -599,11 +555,6 @@ class BaseTaskManager(ABC):
             cleanup_coros = [self._cleanup_task(task_id, mission_id) for task_id in remaining_sessions]
             await asyncio.gather(*cleanup_coros, return_exceptions=True)
 
-        # Await any deferred cleanup tasks
-        if self._cleanup_tasks:
-            await asyncio.gather(*self._cleanup_tasks, return_exceptions=True)
-            self._cleanup_tasks.clear()
-
         logger.info(
             "TaskManager shutdown completed, cancelled: %d, failed: %d",
             len(results) - len(failed_tasks),

{digitalkin-1.0.0.dev2 → digitalkin-1.0.0.dev3}/src/digitalkin/core/task_manager/local_task_manager.py

@@ -62,15 +62,20 @@ class LocalTaskManager(BaseTaskManager):
                 },
             )
 
-            # Execute task using TaskExecutor
+            # Execute task using TaskExecutor; cleanup runs inside the
+            # supervisor's `finally` (folded from former _deferred_cleanup).
+            async def _finalize() -> None:
+                await self._cleanup_task(task_id, mission_id=mission_id)
+
             supervisor_task = await self._executor.execute_task(
                 task_id,
                 mission_id,
                 coro,
                 session,
+                on_finalize=_finalize,
+                stream_drain_timeout=self._stream_drain_timeout,
             )
             self.tasks[task_id] = supervisor_task
-            self._register_auto_cleanup(task_id, mission_id)
 
             logger.info(
                 "Local task created and started: '%s'",