feat(mem_scheduler): add messages logging for stuck tasks monitoring

- Add RunningTaskItem schema with optional messages field in task_schemas.py
- Update dispatcher to pass messages when creating RunningTaskItem instances
- Enhance dispatcher_monitor to log messages info for stuck tasks (count + first 3 messages)
- Add comprehensive unit tests for new messages functionality
- Fix existing test assertions to handle dispatcher's message grouping logic

This improvement provides better debugging visibility for stuck tasks by including
the actual message content and count in monitoring logs.

Author: tangg555

src/memos/mem_scheduler/base_scheduler.py

@@ -59,7 +59,7 @@ def __init__(self, config: BaseSchedulerConfig):
         self.enable_activation_memory = self.config.get("enable_activation_memory", False)
         self.act_mem_dump_path = self.config.get("act_mem_dump_path", DEFAULT_ACT_MEM_DUMP_PATH)
         self.search_method = TreeTextMemory_SEARCH_METHOD
-        self.enable_parallel_dispatch = self.config.get("enable_parallel_dispatch", False)
+        self.enable_parallel_dispatch = self.config.get("enable_parallel_dispatch", True)
         self.thread_pool_max_workers = self.config.get(
             "thread_pool_max_workers", DEFAULT_THREAD_POOL_MAX_WORKERS
         )

src/memos/mem_scheduler/general_modules/dispatcher.py

@@ -10,6 +10,7 @@
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.general_modules.task_threads import ThreadManager
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.task_schemas import RunningTaskItem
 
 
 logger = get_logger(__name__)
@@ -28,7 +29,7 @@ class SchedulerDispatcher(BaseSchedulerModule):
     - Thread race competition for parallel task execution
     """
 
-    def __init__(self, max_workers=30, enable_parallel_dispatch=False, config=None):
+    def __init__(self, max_workers=30, enable_parallel_dispatch=True, config=None):
         super().__init__()
         self.config = config
 
@@ -58,6 +59,68 @@ def __init__(self, max_workers=30, enable_parallel_dispatch=False, config=None):
         # Thread race module for competitive task execution
         self.thread_manager = ThreadManager(thread_pool_executor=self.dispatcher_executor)
 
+        # Task tracking for monitoring
+        self._running_tasks: dict[str, RunningTaskItem] = {}
+        self._task_lock = threading.Lock()
+
+    def _create_task_wrapper(self, handler: Callable, task_item: RunningTaskItem):
+        """
+        Create a wrapper around the handler to track task execution and capture results.
+
+        Args:
+            handler: The original handler function
+            task_item: The RunningTaskItem to track
+
+        Returns:
+            Wrapped handler function that captures results and logs completion
+        """
+
+        def wrapped_handler(messages: list[ScheduleMessageItem]):
+            try:
+                # Execute the original handler
+                result = handler(messages)
+
+                # Mark task as completed and remove from tracking
+                with self._task_lock:
+                    if task_item.item_id in self._running_tasks:
+                        task_item.mark_completed(result)
+                        del self._running_tasks[task_item.item_id]
+
+                logger.info(f"Task completed: {task_item.get_execution_info()}")
+                return result
+
+            except Exception as e:
+                # Mark task as failed and remove from tracking
+                with self._task_lock:
+                    if task_item.item_id in self._running_tasks:
+                        task_item.mark_failed(str(e))
+                        del self._running_tasks[task_item.item_id]
+
+                logger.error(f"Task failed: {task_item.get_execution_info()}, Error: {e}")
+                raise
+
+        return wrapped_handler
+
+    def get_running_tasks(self) -> dict[str, RunningTaskItem]:
+        """
+        Get a copy of currently running tasks.
+
+        Returns:
+            Dictionary of running tasks keyed by task ID
+        """
+        with self._task_lock:
+            return self._running_tasks.copy()
+
+    def get_running_task_count(self) -> int:
+        """
+        Get the count of currently running tasks.
+
+        Returns:
+            Number of running tasks
+        """
+        with self._task_lock:
+            return len(self._running_tasks)
+
     def register_handler(self, label: str, handler: Callable[[list[ScheduleMessageItem]], None]):
         """
         Register a handler function for a specific message label.
@@ -126,7 +189,7 @@ def unregister_handlers(self, labels: list[str]) -> dict[str, bool]:
     def _default_message_handler(self, messages: list[ScheduleMessageItem]) -> None:
         logger.debug(f"Using _default_message_handler to deal with messages: {messages}")
 
-    def group_messages_by_user_and_cube(
+    def _group_messages_by_user_and_mem_cube(
         self, messages: list[ScheduleMessageItem]
     ) -> dict[str, dict[str, list[ScheduleMessageItem]]]:
         """
@@ -176,25 +239,51 @@ def dispatch(self, msg_list: list[ScheduleMessageItem]):
             logger.debug("Received empty message list, skipping dispatch")
             return
 
-        # Group messages by their labels, and organize messages by label
-        label_groups = defaultdict(list)
-        for message in msg_list:
-            label_groups[message.label].append(message)
-
-        # Process each label group
-        for label, msgs in label_groups.items():
-            handler = self.handlers.get(label, self._default_message_handler)
-
-            # dispatch to different handler
-            logger.debug(f"Dispatch {len(msgs)} message(s) to {label} handler.")
-            if self.enable_parallel_dispatch and self.dispatcher_executor is not None:
-                # Capture variables in lambda to avoid loop variable issues
-                future = self.dispatcher_executor.submit(handler, msgs)
-                self._futures.add(future)
-                future.add_done_callback(self._handle_future_result)
-                logger.info(f"Dispatched {len(msgs)} message(s) as future task")
-            else:
-                handler(msgs)
+        # Group messages by user_id and mem_cube_id first
+        user_cube_groups = self._group_messages_by_user_and_mem_cube(msg_list)
+
+        # Process each user and mem_cube combination
+        for user_id, cube_groups in user_cube_groups.items():
+            for mem_cube_id, user_cube_msgs in cube_groups.items():
+                # Group messages by their labels within each user/mem_cube combination
+                label_groups = defaultdict(list)
+                for message in user_cube_msgs:
+                    label_groups[message.label].append(message)
+
+                # Process each label group within this user/mem_cube combination
+                for label, msgs in label_groups.items():
+                    handler = self.handlers.get(label, self._default_message_handler)
+
+                    # Create task tracking item for this dispatch
+                    task_item = RunningTaskItem(
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
+                        task_info=f"Processing {len(msgs)} message(s) with label '{label}' for user {user_id} and mem_cube {mem_cube_id}",
+                        task_name=f"{label}_handler",
+                        messages=msgs,
+                    )
+
+                    # Add to running tasks
+                    with self._task_lock:
+                        self._running_tasks[task_item.item_id] = task_item
+
+                    # Create wrapped handler for task tracking
+                    wrapped_handler = self._create_task_wrapper(handler, task_item)
+
+                    # dispatch to different handler
+                    logger.debug(
+                        f"Dispatch {len(msgs)} message(s) to {label} handler for user {user_id} and mem_cube {mem_cube_id}."
+                    )
+                    logger.info(f"Task started: {task_item.get_execution_info()}")
+
+                    if self.enable_parallel_dispatch and self.dispatcher_executor is not None:
+                        # Capture variables in lambda to avoid loop variable issues
+                        future = self.dispatcher_executor.submit(wrapped_handler, msgs)
+                        self._futures.add(future)
+                        future.add_done_callback(self._handle_future_result)
+                        logger.info(f"Dispatched {len(msgs)} message(s) as future task")
+                    else:
+                        wrapped_handler(msgs)
 
     def join(self, timeout: float | None = None) -> bool:
         """Wait for all dispatched tasks to complete.

src/memos/mem_scheduler/monitors/dispatcher_monitor.py

@@ -9,6 +9,11 @@
 from memos.log import get_logger
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.general_modules.dispatcher import SchedulerDispatcher
+from memos.mem_scheduler.schemas.general_schemas import (
+    DEFAULT_DISPATCHER_MONITOR_CHECK_INTERVAL,
+    DEFAULT_DISPATCHER_MONITOR_MAX_FAILURES,
+    DEFAULT_STUCK_THREAD_TOLERANCE,
+)
 
 
 logger = get_logger(__name__)
@@ -21,8 +26,12 @@ def __init__(self, config: BaseSchedulerConfig):
         super().__init__()
         self.config: BaseSchedulerConfig = config
 
-        self.check_interval = self.config.get("dispatcher_monitor_check_interval", 300)
-        self.max_failures = self.config.get("dispatcher_monitor_max_failures", 2)
+        self.check_interval = self.config.get(
+            "dispatcher_monitor_check_interval", DEFAULT_DISPATCHER_MONITOR_CHECK_INTERVAL
+        )
+        self.max_failures = self.config.get(
+            "dispatcher_monitor_max_failures", DEFAULT_DISPATCHER_MONITOR_MAX_FAILURES
+        )
 
         # Registry of monitored thread pools
         self._pools: dict[str, dict] = {}
@@ -189,22 +198,77 @@ def _check_pools_health(self) -> None:
             ):
                 self._restart_pool(name, pool_info)
 
-    def _check_pool_health(self, pool_info: dict, stuck_max_interval=4) -> tuple[bool, str]:
+    def _check_pool_health(
+        self, pool_info: dict, stuck_max_interval=4, stuck_thread_tolerance=None
+    ) -> tuple[bool, str]:
         """
-        Check health of a single thread pool.
+        Check health of a single thread pool with enhanced task tracking.
 
         Args:
             pool_info: Dictionary containing pool configuration
+            stuck_max_interval: Maximum intervals before considering pool stuck
+            stuck_thread_tolerance: Maximum number of stuck threads to tolerate before restarting pool
 
         Returns:
             Tuple: (is_healthy, reason) where reason explains failure if not healthy
         """
+        if stuck_thread_tolerance is None:
+            stuck_thread_tolerance = DEFAULT_STUCK_THREAD_TOLERANCE
+
         executor = pool_info["executor"]
 
         # Check if executor is shutdown
         if executor._shutdown:  # pylint: disable=protected-access
             return False, "Executor is shutdown"
 
+        # Enhanced health check using dispatcher task tracking
+        stuck_tasks = []
+        if self.dispatcher:
+            running_tasks = self.dispatcher.get_running_tasks()
+            running_count = self.dispatcher.get_running_task_count()
+
+            # Log detailed task information
+            if running_tasks:
+                logger.debug(f"Currently running {running_count} tasks:")
+                for _task_id, task in running_tasks.items():
+                    logger.debug(f"  - {task.get_execution_info()}")
+            else:
+                logger.debug("No tasks currently running")
+
+            # Check for stuck tasks (running longer than expected)
+            for task in running_tasks.values():
+                if task.duration_seconds and task.duration_seconds > (
+                    self.check_interval * stuck_max_interval
+                ):
+                    stuck_tasks.append(task)
+
+            # Always log stuck tasks if any exist
+            if stuck_tasks:
+                logger.warning(f"Found {len(stuck_tasks)} potentially stuck tasks:")
+                for task in stuck_tasks:
+                    task_info = task.get_execution_info()
+                    messages_info = ""
+                    if task.messages:
+                        messages_info = f", Messages: {len(task.messages)} items - {[str(msg) for msg in task.messages[:3]]}"
+                        if len(task.messages) > 3:
+                            messages_info += f" ... and {len(task.messages) - 3} more"
+                    logger.warning(f"  - Stuck task: {task_info}{messages_info}")
+
+                # Check if stuck task count exceeds tolerance
+                # If thread pool size is smaller, use the smaller value as threshold
+                max_workers = pool_info.get("max_workers", 0)
+                effective_tolerance = (
+                    min(stuck_thread_tolerance, max_workers)
+                    if max_workers > 0
+                    else stuck_thread_tolerance
+                )
+
+                if len(stuck_tasks) >= effective_tolerance:
+                    return (
+                        False,
+                        f"Found {len(stuck_tasks)} stuck tasks (tolerance: {effective_tolerance})",
+                    )
+
         # Check thread activity
         active_threads = sum(
             1
@@ -216,13 +280,24 @@ def _check_pool_health(self, pool_info: dict, stuck_max_interval=4) -> tuple[boo
         if active_threads == 0 and pool_info["max_workers"] > 0:
             return False, "No active worker threads"
 
-        # Check if threads are stuck (no activity for 2 intervals)
+        # Check if threads are stuck (no activity for specified intervals)
         time_delta = (datetime.utcnow() - pool_info["last_active"]).total_seconds()
         if time_delta >= self.check_interval * stuck_max_interval:
-            return False, "No recent activity"
+            return False, f"No recent activity for {time_delta:.1f} seconds"
 
         # If we got here, pool appears healthy
         pool_info["last_active"] = datetime.utcnow()
+
+        # Log health status with comprehensive information
+        if self.dispatcher:
+            task_count = self.dispatcher.get_running_task_count()
+            max_workers = pool_info.get("max_workers", 0)
+            stuck_count = len(stuck_tasks)
+            logger.info(
+                f"Pool health check passed - {active_threads} active threads, "
+                f"{task_count} running tasks, pool size: {max_workers}, stuck tasks: {stuck_count}"
+            )
+
         return True, ""
 
     def _restart_pool(self, name: str, pool_info: dict) -> None:

src/memos/mem_scheduler/schemas/general_schemas.py

@@ -17,8 +17,12 @@
 DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT = 30
 DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT = 20
 DEFAULT_ACT_MEM_DUMP_PATH = f"{BASE_DIR}/outputs/mem_scheduler/mem_cube_scheduler_test.kv_cache"
-DEFAULT_THREAD_POOL_MAX_WORKERS = 10
+DEFAULT_THREAD_POOL_MAX_WORKERS = 30
 DEFAULT_CONSUME_INTERVAL_SECONDS = 0.05
+DEFAULT_DISPATCHER_MONITOR_CHECK_INTERVAL = 300
+DEFAULT_DISPATCHER_MONITOR_MAX_FAILURES = 2
+DEFAULT_STUCK_THREAD_TOLERANCE = 10
+
 NOT_INITIALIZED = -1
 
 

src/memos/mem_scheduler/schemas/task_schemas.py

@@ -0,0 +1,67 @@
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, computed_field
+
+from memos.log import get_logger
+from memos.mem_scheduler.general_modules.misc import DictConversionMixin
+
+
+logger = get_logger(__name__)
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
+
+
+# ============== Running Tasks ==============
+class RunningTaskItem(BaseModel, DictConversionMixin):
+    """Data class for tracking running tasks in SchedulerDispatcher."""
+
+    item_id: str = Field(
+        description="Unique identifier for the task item", default_factory=lambda: str(uuid4())
+    )
+    user_id: str = Field(..., description="Required user identifier", min_length=1)
+    mem_cube_id: str = Field(..., description="Required memory cube identifier", min_length=1)
+    task_info: str = Field(..., description="Information about the task being executed")
+    task_name: str = Field(..., description="Name/type of the task handler")
+    start_time: datetime = Field(description="Task start time", default_factory=datetime.utcnow)
+    end_time: datetime | None = Field(default=None, description="Task completion time")
+    status: str = Field(default="running", description="Task status: running, completed, failed")
+    result: Any | None = Field(default=None, description="Task execution result")
+    error_message: str | None = Field(default=None, description="Error message if task failed")
+    messages: list[Any] | None = Field(
+        default=None, description="List of messages being processed by this task"
+    )
+
+    def mark_completed(self, result: Any | None = None) -> None:
+        """Mark task as completed with optional result."""
+        self.end_time = datetime.utcnow()
+        self.status = "completed"
+        self.result = result
+
+    def mark_failed(self, error_message: str) -> None:
+        """Mark task as failed with error message."""
+        self.end_time = datetime.utcnow()
+        self.status = "failed"
+        self.error_message = error_message
+
+    @computed_field
+    @property
+    def duration_seconds(self) -> float | None:
+        """Calculate task duration in seconds."""
+        if self.end_time:
+            return (self.end_time - self.start_time).total_seconds()
+        return None
+
+    def get_execution_info(self) -> str:
+        """Get formatted execution information for logging."""
+        duration = self.duration_seconds
+        duration_str = f"{duration:.2f}s" if duration else "ongoing"
+
+        return (
+            f"Task {self.task_name} (ID: {self.item_id[:8]}) "
+            f"for user {self.user_id}, cube {self.mem_cube_id} - "
+            f"Status: {self.status}, Duration: {duration_str}"
+        )