ws-bom-robot-app 0.0.63__py3-none-any.whl → 0.0.103__py3-none-any.whl

ws_bom_robot_app/task_manager.py

@@ -1,7 +1,7 @@
 from collections import deque
 import inspect
 from math import floor
-import asyncio, os
+import asyncio, os, traceback
 from datetime import datetime, timedelta
 from enum import Enum
 from typing import Annotated, Coroutine, Literal, TypeVar, Optional, Dict, Union, Any, Callable
@@ -15,6 +15,8 @@ from ws_bom_robot_app.util import _log
 from sqlalchemy import create_engine, Column, String, JSON, DateTime, Enum
 from sqlalchemy.orm import sessionmaker, registry
 from abc import ABC, abstractmethod
+from ws_bom_robot_app.subprocess_runner import _start_subprocess_for_coroutine, _recv_from_connection_async, _pickler
+from ws_bom_robot_app.config import config
 
 T = TypeVar('T')
 
@@ -58,7 +60,8 @@ class TaskMetaData(BaseModel):
             )
     source: Optional[str] = None
     pid: Optional[int] = None
-    extra: Optional[dict[str,str]] = None
+    pid_child: Optional[int] = None
+    extra: Optional[dict[str,Union[str,int,bool]]] = None
 
 class TaskStatus(IdentifiableEntity):
     type: Optional[str] = None
@@ -66,19 +69,21 @@ class TaskStatus(IdentifiableEntity):
     result: Optional[T] = None
     metadata: TaskMetaData = None
     error: Optional[str] = None
+    retry: int = 0
     model_config = ConfigDict(
         arbitrary_types_allowed=True
     )
 
 class TaskEntry(IdentifiableEntity):
-    task: Annotated[asyncio.Task, Field(default=None, validate_default=False)] = None
-    coroutine: Coroutine = None
+    task: Annotated[Union[asyncio.Task, Callable], Field(default=None, validate_default=False)] = None
+    coroutine: Any = None
     headers: TaskHeader | None = None
     status: Union[TaskStatus, None] = None
-    def _get_coroutine_name(self, coroutine: asyncio.coroutines) -> str:
-        if inspect.iscoroutine(coroutine):
-            return coroutine.cr_code.co_name
-        return "<unknown>"
+    queue: Literal["slow", "fast"] | None = "slow"
+    def _get_coroutine_name(self, func: Any) -> str:
+        if inspect.iscoroutine(func):
+            return func.cr_code.co_name
+        return func.__qualname__ if callable(func) else str(func)
     def __init__(self, **data):
         def _metadata_extra(data: str) -> dict[str,str] | None:
             if data:
@@ -115,7 +120,8 @@ class TaskEntry(IdentifiableEntity):
 class TaskStatistics(BaseModel):
     class TaskStatisticExecutionInfo(BaseModel):
         retention_days: float = config.robot_task_retention_days
-        max_concurrent: int
+        max_parallelism: int
+        slot_available: dict[str,int]
         pid: int = os.getpid()
         running: list[TaskStatus]
         slowest: list
@@ -134,14 +140,29 @@ class TaskStatistics(BaseModel):
 
 #region interface
 class TaskManagerStrategy(ABC):
-    def __init__(self, max_concurrent_tasks: int = max(1,floor(config.robot_task_max_total_parallelism / config.runtime_options().number_of_workers))):
-        self.max_concurrent_tasks = max_concurrent_tasks
-        self.semaphore = asyncio.Semaphore(self.max_concurrent_tasks)
+    def __init__(self, max_concurrent_tasks: Optional[int] = None):
+        if max_concurrent_tasks is None:
+            workers = config.runtime_options().number_of_workers
+            max_concurrent_tasks = max(1, floor(config.robot_task_max_total_parallelism / max(1, workers)))
+        self.max_parallelism = max_concurrent_tasks
+        self.semaphore = {"slow": asyncio.Semaphore(max_concurrent_tasks), "fast": asyncio.Semaphore(max_concurrent_tasks*2)}
         self.running_tasks = dict[str, TaskEntry]()
         self.loop = asyncio.get_event_loop()
 
     @abstractmethod
-    def create_task(self, coroutine, headers: TaskHeader | None = None) -> IdentifiableEntity:
+    def create_task(self, coroutine, headers: TaskHeader | None = None, queue: Literal["slow", "fast"] | None = "slow") -> IdentifiableEntity:
+        """Create a new task.
+        Args:
+            coroutine (_type_): coroutine or callable to be executed.
+            headers (TaskHeader | None, optional): for api call, http headers to include with the task. Defaults to None.
+        Returns:
+            IdentifiableEntity: The created task id.
+        Usage:
+            from ws_bom_robot_app.task_manager import task_manager
+            task_manager.create_task(my_coroutine, headers=my_headers) -> coroutine executed in-process
+            task_manager.create_task(lambda: my_coroutine, headers=my_headers) -> callable using subprocess
+            task_manager.create_task(lambda: my_coroutine, headers=my_headers, queue="fast") -> callable using subprocess with "fast" queue
+        """
         pass
 
     @abstractmethod
@@ -172,49 +193,173 @@ class TaskManagerStrategy(ABC):
     def task_cleanup_rule(self, task: TaskEntry) -> bool:
         return task.status.metadata.start_at and datetime.fromisoformat(task.status.metadata.start_at) < datetime.now() - timedelta(days=config.robot_task_retention_days)
 
-    def task_done_callback(self, task_entry: TaskEntry, headers: TaskHeader | None = None) -> Callable:
-        def callback(task: asyncio.Task):
+    def _update_task_by_event(self, task_entry: TaskEntry, status: str, output: Any) -> None:
+        if status == "completed":
+          task_entry.status.status = "completed"
+          task_entry.status.result = output
+        elif status == "failure":
+          task_entry.status.status = "failure"
+          task_entry.status.error = str(output)
+          _log.error(f"Task {task_entry.id} failed with error: {output}")
+        else:
+          task_entry.status.metadata.end_at = str(datetime.now().isoformat())
+          #strategy-specific behavior
+          self.update_task_status(task_entry)
+          #remove from running tasks
+          if task_entry.id in self.running_tasks:
+              del self.running_tasks[task_entry.id]
+          #notify webhooks: a task has completed or failed, if failed with retry policy the task remains in pending state, and will not be notified until complete/failure
+          if task_entry.status.status in ["completed","failure"]:
+            if task_entry.headers and task_entry.headers.x_ws_bom_webhooks:
+              try:
+                  asyncio.create_task(
+                      WebhookNotifier().notify_webhook(task_entry.status, task_entry.headers.x_ws_bom_webhooks)
+                  )
+              except Exception as e:
+                  _log.error(f"Failed to schedule webhook notification for task {task_entry.id}: {e}")
+
+    def task_done_callback(self, task_entry: TaskEntry) -> Callable:
+        def callback(task: asyncio.Task, context: Any | None = None):
             try:
-                result = task_entry.task.result()
-                task_entry.status.status = "completed"
-                task_entry.status.result = result
+                result = task.result()
+                self._update_task_by_event(task_entry, "completed", result)
             except Exception as e:
-                task_entry.status.status = "failure"
-                task_entry.status.error = str(e)
+                self._update_task_by_event(task_entry, "failure", e)
             finally:
-                task_entry.status.metadata.end_at = str(datetime.now().isoformat())
-                #strategy-specific behavior
-                self.update_task_status(task_entry)
-                #remove from running tasks
-                del self.running_tasks[task_entry.id]
-                #notify webhooks
-                if headers and headers.x_ws_bom_webhooks:
-                  asyncio.create_task(
-                      WebhookNotifier().notify_webhook(task_entry.status,headers.x_ws_bom_webhooks)
-                      )
+                self._update_task_by_event(task_entry, "callback", None)
         return callback
 
-    def create_task_entry(self, coroutine: asyncio.coroutines, headers: TaskHeader | None = None) -> TaskEntry:
+    def create_task_entry(self, coroutine_or_callable: Any, headers: TaskHeader | None = None, queue: Literal["slow", "fast"] | None = "slow") -> TaskEntry:
+        """Create a new task entry.
+
+        Args:
+            coroutine_or_callable (Any): The coroutine or callable to be executed.
+            headers (TaskHeader | None, optional): Headers to include with the task. Defaults to None.
+        Raises:
+            TypeError: If the input is not a coroutine or callable.
+        Returns:
+            TaskEntry: The created task entry.
+        """
         _id = headers and headers.x_ws_bom_msg_id or str(uuid4())
-        task = TaskEntry(
+        # Detect coroutine object
+        if inspect.iscoroutine(coroutine_or_callable):
+            can_use_subprocess = False
+        elif callable(coroutine_or_callable):
+            can_use_subprocess = True
+        else:
+            raise TypeError(
+                f"Expected coroutine object or callable, got {type(coroutine_or_callable)}"
+            )
+        task_entry = TaskEntry(
             id=_id,
-            coroutine=coroutine,
-            headers=headers)
-        self.loop.create_task(self._run_task_with_semaphore(task)) # run the task
-        return task
+            coroutine=coroutine_or_callable,
+            headers=headers,
+            queue=queue
+        )
+        # Store hint for subprocess capability
+        task_entry.status.metadata.extra = task_entry.status.metadata.extra or {}
+        task_entry.status.metadata.extra["can_use_subprocess"] = can_use_subprocess
+        try:
+          asyncio.create_task(self._run_task_with_semaphore(task_entry)) # run the task
+        except Exception as e:
+          _log.error(f"Error occurred while creating task {task_entry.id}: {e}")
+        return task_entry
 
     async def _run_task_with_semaphore(self, task_entry: TaskEntry):
         """Run a task with semaphore control to limit concurrency."""
-        async with self.semaphore:
+        async with self.semaphore[task_entry.queue]:
           await self._execute_task(task_entry)
 
+    async def _monitor_subprocess(self, task_entry: TaskEntry, proc, conn):
+        try:
+            # Wait for the worker to send bytes (this blocks, so run via executor wrapper)
+            data_bytes = await _recv_from_connection_async(conn)
+            # unpickle bytes to get payload
+            try:
+                payload = _pickler.loads(data_bytes)
+            except Exception:
+                # fallback if pickler fails
+                payload = ("err", {"error": "Failed to unpickle subprocess result"})
+            if isinstance(payload, tuple) and payload[0] == "ok":
+                result = payload[1]
+                # write results into task_entry
+                self._update_task_by_event(task_entry, "completed", result)
+            else:
+                # error
+                err_info = payload[1]["error"] if isinstance(payload, tuple) else str(payload)
+                self._update_task_by_event(task_entry, "failure", err_info) # give up, no retry
+        except Exception:
+            # maybe subprocess is no more alive / killed due to memory pressure
+            if task_entry.status.retry < config.robot_task_mp_max_retries:
+                task_entry.status.retry += 1
+                _log.warning(f"Task {task_entry.id} failure, retrying {task_entry.status.retry}...")
+                async def delayed_retry():
+                    _delay = config.robot_task_mp_retry_delay # help to backpressure when overloaded
+                    if self.semaphore[task_entry.queue]._value > 0:  # free semaphore slots available
+                        _delay = 5  # small/no delay if retry can run immediately
+                    await asyncio.sleep(_delay)  # delay in seconds
+                    await self._run_task_with_semaphore(task_entry)
+                asyncio.create_task(delayed_retry())
+                # semaphore is released, so new task can be executed
+                return
+            else:
+                self._update_task_by_event(task_entry, "failure", "subprocess monitor error: failed to receive data from connection")
+        finally:
+            # ensure process termination / cleanup
+            try:
+                conn.close()
+            except Exception:
+                pass
+            try:
+              if proc.is_alive():
+                  proc.terminate()
+              proc.join(timeout=1)
+            except Exception:
+                pass
+            # callback
+            self._update_task_by_event(task_entry, "callback", None)
+
     async def _execute_task(self, task_entry: TaskEntry):
-        """Execute a task and handle its lifecycle."""
+        """
+        Execute the task. Try to run it inside a subprocess (if serializable).
+        If subprocess is used, we create a monitor asyncio.Task that waits for the subprocess result
+        and then calls the same task_done_callback to finalize and persist state.
+        If subprocess cannot be used, fall back to in-process behavior.
+        """
         self.running_tasks[task_entry.id]=task_entry
         task_entry.status.metadata.start_at = str(datetime.now().isoformat())
-        task_entry.task = asyncio.create_task(task_entry.coroutine)
-        task_entry.task.add_done_callback(self.task_done_callback(task_entry, task_entry.headers))
-        await task_entry.task
+        # try to spawn subprocess (non-blocking)
+        can_use_subprocess = task_entry.status.metadata.extra.get("can_use_subprocess", False)
+        if config.robot_task_mp_enable and can_use_subprocess:
+            proc, conn, used_subprocess = _start_subprocess_for_coroutine(task_entry.coroutine)
+            if used_subprocess and proc is not None and conn is not None:
+                # monitor subprocess asynchronously
+                task_entry.status.status = "pending"
+                task_entry.status.metadata.pid_child = proc.pid
+                _log.info(f"Task {task_entry.id} started in subprocess (pid={proc.pid})")
+                # await monitor process, then return: important to acquire semaphore
+                await self._monitor_subprocess(task_entry, proc, conn)
+                return
+        # default fallback (in-process)
+        try:
+            async def _callable_to_coroutine(func: Any) -> Any:
+                if callable(func) and not inspect.iscoroutine(func):
+                    result = func()
+                    if inspect.iscoroutine(result):
+                        return await result
+                    return result
+                elif inspect.iscoroutine(func):
+                    return await func
+                return func
+            task_entry.task = asyncio.create_task(_callable_to_coroutine(task_entry.coroutine))
+            task_entry.task.add_done_callback(self.task_done_callback(task_entry))
+            _log.info(f"Starting task {task_entry.id} in-process with coroutine {task_entry._get_coroutine_name(task_entry.coroutine)}")
+            await task_entry.task
+        except Exception as e:
+            _error = f"Error occurred while executing task {task_entry.id}: {e}"
+            _log.error(_error)
+            self._update_task_by_event(task_entry, "failure", _error)
+            self._update_task_by_event(task_entry, "callback", None)
 
     def running_task(self):
         return self.running_tasks.values()
@@ -250,7 +395,8 @@ class TaskManagerStrategy(ABC):
             ),
             exec_info=TaskStatistics.TaskStatisticExecutionInfo(
                 retention_days=config.robot_task_retention_days,
-                max_concurrent=self.max_concurrent_tasks,
+                max_parallelism=self.max_parallelism,
+                slot_available={queue: self.semaphore[queue]._value for queue in self.semaphore},
                 running=[task.status for task in self.running_task()],
                 slowest=_slowest
             )
@@ -258,14 +404,14 @@ class TaskManagerStrategy(ABC):
 
 #endregion
 
-#memory implementation
+#region memory implementation
 class MemoryTaskManagerStrategy(TaskManagerStrategy):
-    def __init__(self):
-        super().__init__()
+    def __init__(self, max_concurrent_tasks: Optional[int] = None):
+        super().__init__(max_concurrent_tasks)
         self.tasks: Dict[str, TaskEntry] = {}
 
-    def create_task(self, coroutine: asyncio.coroutines, headers: TaskHeader | None = None) -> IdentifiableEntity:
-        task = self.create_task_entry(coroutine, headers)
+    def create_task(self, coroutine: Any, headers: TaskHeader | None = None, queue: Literal["slow", "fast"] | None = "slow") -> IdentifiableEntity:
+        task = self.create_task_entry(coroutine, headers, queue)
         self.tasks[task.id] = task
         return IdentifiableEntity(id=task.id)
 
@@ -292,7 +438,7 @@ class MemoryTaskManagerStrategy(TaskManagerStrategy):
 
 #endregion
 
-#db implementation
+#region db implementation
 Base = registry().generate_base()
 class TaskEntryModel(Base):
     __tablename__ = "entry"
@@ -302,8 +448,8 @@ class TaskEntryModel(Base):
         arbitrary_types_allowed=True
     )
 class DatabaseTaskManagerStrategy(TaskManagerStrategy):
-    def __init__(self, db_url: str = "sqlite:///.data/db/tasks.sqlite"):
-        super().__init__()
+    def __init__(self, db_url: str = f"sqlite:///{config.robot_data_folder}/db/tasks.sqlite", max_concurrent_tasks: Optional[int] = None):
+        super().__init__(max_concurrent_tasks)
         self.engine = create_engine(db_url)
         self.Session = sessionmaker(bind=self.engine)
         Base.metadata.create_all(self.engine)
@@ -353,11 +499,16 @@ class DatabaseTaskManagerStrategy(TaskManagerStrategy):
 #region global
 def __get_taskmanager_strategy() -> TaskManagerStrategy:
     """ Factory function to get the appropriate task manager strategy based on the runtime configuration."""
-    if config.runtime_options().is_multi_process:
-        return DatabaseTaskManagerStrategy()
-    return MemoryTaskManagerStrategy()
+    if config.robot_task_strategy == 'memory':
+        return MemoryTaskManagerStrategy()
+    return DatabaseTaskManagerStrategy()
 task_manager = __get_taskmanager_strategy()
 _log.info(f"Task manager strategy: {task_manager.__class__.__name__}")
+
+def task_cleanup():
+    _log.info("Cleaning up tasks...")
+    task_manager.cleanup_task()
+    _log.info("Task cleanup complete.")
 #endregion
 
 #region api

ws_bom_robot_app/util.py

@@ -22,6 +22,12 @@ def logger_instance(name: str) -> logging.Logger:
 _log: logging.Logger = locals().get("_loc", logger_instance(__name__))
 #endregion
 
+#region task
+def is_app_subprocess():
+    """Check if we're running a task in a subprocess."""
+    return os.environ.get('IS_ROBOT_APP_SUBPROCESS', '').lower() == 'true'
+#endregion
+
 #region cache
 _cache = {}
 _cache_timestamps = {}