address the incompatible issue of local scheduler

Author: tangg555

src/memos/mem_scheduler/base_scheduler.py

@@ -1009,14 +1009,24 @@ def _monitor_loop(self):
                 q_sizes = self.memos_message_queue.qsize()
 
                 for stream_key, queue_length in q_sizes.items():
-                    # Expected format: "memos:stream:{user_id}:{mem_cube_id}" or "{user_id}"
+                    # Skip aggregate keys like 'total_size'
+                    if stream_key == "total_size":
+                        continue
+
+                    # Key format: ...:{user_id}:{mem_cube_id}:{task_label}
+                    # We want to extract user_id, which is the 3rd component from the end.
                     parts = stream_key.split(":")
                     if len(parts) >= 3:
-                        user_id = parts[2]
-                        self.metrics.update_queue_length(queue_length, user_id)
-                    elif not self.use_redis_queue:  # local queue
-                        user_id = stream_key
+                        user_id = parts[-3]
                         self.metrics.update_queue_length(queue_length, user_id)
+                    else:
+                        # Fallback for unexpected key formats (e.g. legacy or testing)
+                        # Try to use the key itself if it looks like a user_id (no colons)
+                        # or just log a warning?
+                        # For now, let's assume if it's not total_size and short, it might be a direct user_id key
+                        # (though that shouldn't happen with current queue implementations)
+                        if ":" not in stream_key:
+                            self.metrics.update_queue_length(queue_length, stream_key)
 
             except Exception as e:
                 logger.error(f"Error in metrics monitor loop: {e}", exc_info=True)

src/memos/mem_scheduler/task_schedule_modules/local_queue.py

@@ -62,7 +62,7 @@ def put(
             Exception: Any underlying error during queue.put() operation.
         """
         stream_key = self.get_stream_key(
-            user_id=message.user_id, mem_cube_id=message.mem_cube_id, task_label=message.task_label
+            user_id=message.user_id, mem_cube_id=message.mem_cube_id, task_label=message.label
         )
 
         message.stream_key = stream_key
@@ -108,35 +108,95 @@ def get(
         )
         return res
 
-    def get_nowait(self, batch_size: int | None = None) -> list[ScheduleMessageItem]:
+    def get_nowait(
+        self, stream_key: str, batch_size: int | None = None
+    ) -> list[ScheduleMessageItem]:
         """
-        Non-blocking version of get(). Equivalent to get(block=False, batch_size=batch_size).
+        Non-blocking version of get(). Equivalent to get(stream_key, block=False, batch_size=batch_size).
 
         Returns immediately with available messages or an empty list if queue is empty.
 
         Args:
+            stream_key (str): The stream/queue identifier.
             batch_size (int | None): Number of messages to retrieve in a batch.
                                    If None, retrieves one message.
 
         Returns:
             List[ScheduleMessageItem]: Retrieved messages or empty list if queue is empty.
         """
-        logger.debug(f"get_nowait() called with batch_size: {batch_size}")
-        return self.get(block=False, batch_size=batch_size)
+        logger.debug(f"get_nowait() called for {stream_key} with batch_size: {batch_size}")
+        return self.get(stream_key=stream_key, block=False, batch_size=batch_size)
+
+    def get_messages(self, batch_size: int) -> list[ScheduleMessageItem]:
+        """
+        Get messages from all streams in round-robin or sequential fashion.
+        Equivalent to SchedulerRedisQueue.get_messages.
+        """
+        messages = []
+        # Snapshot keys to avoid runtime modification issues
+        stream_keys = list(self.queue_streams.keys())
+
+        # Simple strategy: try to get up to batch_size messages across all streams
+        # We can just iterate and collect.
+
+        # Calculate how many to get per stream to be fair?
+        # Or just greedy? Redis implementation uses a complex logic.
+        # For local, let's keep it simple: just iterate and take what's available (non-blocking)
+
+        for stream_key in stream_keys:
+            if len(messages) >= batch_size:
+                break
+
+            needed = batch_size - len(messages)
+            # Use get_nowait to avoid blocking
+            fetched = self.get_nowait(stream_key=stream_key, batch_size=needed)
+            messages.extend(fetched)
+
+        return messages
 
     def qsize(self) -> dict:
         """
         Return the current size of all internal queues as a dictionary.
 
         Each key is the stream name, and each value is the number of messages in that queue.
+        Also includes 'total_size'.
 
         Returns:
             Dict[str, int]: Mapping from stream name to current queue size.
         """
         sizes = {stream: queue.qsize() for stream, queue in self.queue_streams.items()}
+        total_size = sum(sizes.values())
+        sizes["total_size"] = total_size
         logger.debug(f"Current queue sizes: {sizes}")
         return sizes
 
+    def size(self) -> int:
+        """
+        Get the current size of the queue (total message count).
+        Compatible with SchedulerRedisQueue.
+        """
+        return self.unfinished_tasks
+
+    def empty(self) -> bool:
+        """
+        Check if the queue is empty.
+        Compatible with SchedulerRedisQueue.
+        """
+        return self.size() == 0
+
+    def full(self) -> bool:
+        """
+        Check if the queue is full.
+        Compatible with SchedulerRedisQueue.
+
+        Returns True if all internal queues are full.
+        If there are no queues, returns False.
+        """
+        if not self.queue_streams:
+            return False
+
+        return all(queue.full() for queue in self.queue_streams.values())
+
     def clear(self) -> None:
         for queue in self.queue_streams.values():
             queue.clear()
@@ -151,6 +211,9 @@ def unfinished_tasks(self) -> int:
         Returns:
             int: Sum of all message counts in all internal queues.
         """
-        total = sum(self.qsize().values())
+        # qsize() now includes "total_size", so we need to be careful not to double count if we use qsize() values
+        # But qsize() implementation above sums values from queue_streams, then adds total_size.
+        # So sum(self.queue_streams.values().qsize()) is safer.
+        total = sum(queue.qsize() for queue in self.queue_streams.values())
         logger.debug(f"Total unfinished tasks across all queues: {total}")
         return total

src/memos/mem_scheduler/task_schedule_modules/task_queue.py

@@ -153,28 +153,7 @@ def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageIt
                         )
 
     def get_messages(self, batch_size: int) -> list[ScheduleMessageItem]:
-        if isinstance(self.memos_message_queue, SchedulerRedisQueue):
-            return self.memos_message_queue.get_messages(batch_size=batch_size)
-        stream_keys = self.get_stream_keys()
-
-        if len(stream_keys) == 0:
-            return []
-
-        messages: list[ScheduleMessageItem] = []
-
-        for stream_key in stream_keys:
-            fetched = self.memos_message_queue.get(
-                stream_key=stream_key,
-                block=False,
-                batch_size=batch_size,
-            )
-
-            messages.extend(fetched)
-        if len(messages) > 0:
-            logger.debug(
-                f"Fetched {len(messages)} messages across users with per-user batch_size={batch_size}"
-            )
-        return messages
+        return self.memos_message_queue.get_messages(batch_size=batch_size)
 
     def clear(self):
         self.memos_message_queue.clear()

src/memos/mem_scheduler/utils/status_tracker.py

@@ -13,7 +13,7 @@
 
 class TaskStatusTracker:
     @require_python_package(import_name="redis", install_command="pip install redis")
-    def __init__(self, redis_client: "redis.Redis"):
+    def __init__(self, redis_client: "redis.Redis | None"):
         self.redis = redis_client
 
     def _get_key(self, user_id: str) -> str:
@@ -41,6 +41,9 @@ def task_submitted(
             mem_cube_id: Memory cube identifier
             business_task_id: Optional business-level task ID (one task_id can have multiple item_ids)
         """
+        if not self.redis:
+            return
+
         key = self._get_key(user_id)
         payload = {
             "status": "waiting",
@@ -61,6 +64,9 @@ def task_submitted(
         self.redis.expire(key, timedelta(days=7))
 
     def task_started(self, task_id: str, user_id: str):
+        if not self.redis:
+            return
+
         key = self._get_key(user_id)
         existing_data_json = self.redis.hget(key, task_id)
         if not existing_data_json:
@@ -77,6 +83,9 @@ def task_started(self, task_id: str, user_id: str):
         self.redis.expire(key, timedelta(days=7))
 
     def task_completed(self, task_id: str, user_id: str):
+        if not self.redis:
+            return
+
         key = self._get_key(user_id)
         existing_data_json = self.redis.hget(key, task_id)
         if not existing_data_json:
@@ -91,6 +100,9 @@ def task_completed(self, task_id: str, user_id: str):
         self.redis.expire(key, timedelta(days=7))
 
     def task_failed(self, task_id: str, user_id: str, error_message: str):
+        if not self.redis:
+            return
+
         key = self._get_key(user_id)
         existing_data_json = self.redis.hget(key, task_id)
         if not existing_data_json:
@@ -108,11 +120,17 @@ def task_failed(self, task_id: str, user_id: str, error_message: str):
         self.redis.expire(key, timedelta(days=7))
 
     def get_task_status(self, task_id: str, user_id: str) -> dict | None:
+        if not self.redis:
+            return None
+
         key = self._get_key(user_id)
         data = self.redis.hget(key, task_id)
         return json.loads(data) if data else None
 
     def get_all_tasks_for_user(self, user_id: str) -> dict[str, dict]:
+        if not self.redis:
+            return {}
+
         key = self._get_key(user_id)
         all_tasks = self.redis.hgetall(key)
         return {tid: json.loads(t_data) for tid, t_data in all_tasks.items()}
@@ -132,6 +150,9 @@ def get_task_status_by_business_id(self, business_task_id: str, user_id: str) ->
             - If any item is 'failed' → 'failed'
             Returns None if task_id not found.
         """
+        if not self.redis:
+            return None
+
         # Get all item_ids for this task_id
         task_items_key = self._get_task_items_key(user_id, business_task_id)
         item_ids = self.redis.smembers(task_items_key)
@@ -180,6 +201,9 @@ def get_all_tasks_global(self) -> dict[str, dict[str, dict]]:
         Returns:
             dict: {user_id: {task_id: task_data, ...}, ...}
         """
+        if not self.redis:
+            return {}
+
         all_users_tasks = {}
         cursor: int | str = 0
         while True:

tests/test_local_queue_full.py

@@ -0,0 +1,54 @@
+import unittest
+
+from datetime import datetime, timezone
+
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.task_schedule_modules.local_queue import SchedulerLocalQueue
+
+
+class TestLocalQueueFull(unittest.TestCase):
+    def test_full_behavior(self):
+        # Create a queue with very small maxsize for testing
+        lq = SchedulerLocalQueue(maxsize=1)
+
+        # Initially empty
+        self.assertFalse(lq.full())
+
+        # Add message to stream 1
+        msg1 = ScheduleMessageItem(
+            user_id="u1",
+            mem_cube_id="c1",
+            label="l1",
+            content="m1",
+            timestamp=datetime.now(timezone.utc),
+        )
+        lq.put(msg1)
+
+        # Now stream 1 is full (maxsize=1).
+        # Since it's the only stream, and it's full, lq.full() should be True.
+        self.assertTrue(lq.full())
+
+        # Add message to stream 2
+        msg2 = ScheduleMessageItem(
+            user_id="u2",
+            mem_cube_id="c2",
+            label="l2",
+            content="m2",
+            timestamp=datetime.now(timezone.utc),
+        )
+        lq.put(msg2)
+
+        # Now both stream 1 and stream 2 are full. lq.full() should be True.
+        self.assertTrue(lq.full())
+
+        # Remove message from stream 1
+        stream1_key = lq.get_stream_key("u1", "c1", "l1")
+        lq.get(stream1_key)
+
+        # Now stream 1 is empty, stream 2 is full.
+        # "all streams are full" is False.
+        self.assertFalse(lq.full())
+
+
+if __name__ == "__main__":
+    unittest.main()