fix bugs: address ai review issues

Author: tangg555

evaluation/scripts/locomo/locomo_eval.py

@@ -311,7 +311,7 @@ async def main(frame, version="default", options=None, num_runs=1, max_workers=4
     with open(response_path) as file:
         locomo_responses = json.load(file)
 
-    num_users = 2
+    num_users = 10
     all_grades = {}
 
     total_responses_count = sum(

src/memos/mem_scheduler/monitors/task_schedule_monitor.py

@@ -209,13 +209,12 @@ async def _collect_async() -> dict:
                     return local
 
                 try:
-                    loop = asyncio.get_running_loop()
-                    if loop.is_running():
-                        raise RuntimeError("event loop running")
+                    asyncio.get_running_loop()
+                    loop_running = True
                 except RuntimeError:
-                    loop = None
+                    loop_running = False
 
-                if loop is None:
+                if not loop_running:
                     return asyncio.run(_collect_async())
             except Exception as e:
                 logger.debug(f"Parallel status collection failed, fallback to sequential: {e}")

src/memos/mem_scheduler/task_schedule_modules/dispatcher.py

@@ -93,8 +93,6 @@ def __init__(
         # Task tracking for monitoring
         self._running_tasks: dict[str, RunningTaskItem] = {}
         self._task_lock = threading.Lock()
-        self._completed_tasks = []
-        self.completed_tasks_max_show_size = 10
 
         # Configure shutdown wait behavior from config or default
         self.stop_wait = (
@@ -128,6 +126,10 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                 self.status_tracker.task_started(
                     task_id=task_item.item_id, user_id=task_item.user_id
                 )
+            # Record task as running for monitoring (LocalQueue only)
+            if isinstance(self.memos_message_queue, SchedulerLocalQueue):
+                with self._task_lock:
+                    self._running_tasks[task_item.item_id] = task_item
             try:
                 # --- mark start: record queuing time(now - enqueue_ts)---
                 now = time.time()
@@ -179,14 +181,12 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                             redis_message_id=redis_message_id,
                         )
 
-                # 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]
-                        self._completed_tasks.append(task_item)
-                        if len(self._completed_tasks) > self.completed_tasks_max_show_size:
-                            self._completed_tasks.pop(0)
+                # Mark task as completed and remove from tracking (LocalQueue only)
+                if isinstance(self.memos_message_queue, SchedulerLocalQueue):
+                    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
 
@@ -197,13 +197,12 @@ def wrapped_handler(messages: list[ScheduleMessageItem]):
                     self.status_tracker.task_failed(
                         task_id=task_item.item_id, user_id=task_item.user_id, error_message=str(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]
-                        if len(self._completed_tasks) > self.completed_tasks_max_show_size:
-                            self._completed_tasks.pop(0)
+                # Mark task as failed and remove from tracking (LocalQueue only)
+                if isinstance(self.memos_message_queue, SchedulerLocalQueue):
+                    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
@@ -238,10 +237,20 @@ def get_running_tasks(
                 lambda task: task.user_id == "user123" and task.status == "running"
             )
         """
-        with self._task_lock:
+        # Use lock only for LocalQueue; otherwise read without lock
+        if isinstance(self.memos_message_queue, SchedulerLocalQueue):
+            with self._task_lock:
+                if filter_func is None:
+                    return self._running_tasks.copy()
+
+                return {
+                    task_id: task_item
+                    for task_id, task_item in self._running_tasks.items()
+                    if filter_func(task_item)
+                }
+        else:
             if filter_func is None:
                 return self._running_tasks.copy()
-
             return {
                 task_id: task_item
                 for task_id, task_item in self._running_tasks.items()
@@ -255,7 +264,11 @@ def get_running_task_count(self) -> int:
         Returns:
             Number of running tasks
         """
-        with self._task_lock:
+        # Use lock only for LocalQueue; otherwise read without lock
+        if isinstance(self.memos_message_queue, SchedulerLocalQueue):
+            with self._task_lock:
+                return len(self._running_tasks)
+        else:
             return len(self._running_tasks)
 
     def register_handler(self, label: str, handler: Callable[[list[ScheduleMessageItem]], None]):
@@ -393,10 +406,6 @@ def dispatch(self, msg_list: list[ScheduleMessageItem]):
                         messages=msgs,
                     )
 
-                    # Track running task for status/monitoring
-                    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)
 

src/memos/mem_scheduler/task_schedule_modules/local_queue.py

@@ -61,7 +61,9 @@ def put(
             queue.Full: If the queue is full and block=False or timeout expires.
             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)
+        stream_key = self.get_stream_key(
+            user_id=message.user_id, mem_cube_id=message.mem_cube_id, task_label=message.task_label
+        )
 
         message.stream_key = stream_key
 

src/memos/multi_mem_cube/single_cube.py

@@ -375,8 +375,8 @@ def _search_pref(
         """
         if os.getenv("ENABLE_PREFERENCE_MEMORY", "false").lower() != "true":
             return []
-        print(f"search_req.filter for preference memory: {search_req.filter}")
-        print(f"type of pref_mem: {type(self.naive_mem_cube.pref_mem)}")
+        logger.info(f"search_req.filter for preference memory: {search_req.filter}")
+        logger.info(f"type of pref_mem: {type(self.naive_mem_cube.pref_mem)}")
         try:
             results = self.naive_mem_cube.pref_mem.search(
                 query=search_req.query,

src/memos/vec_dbs/milvus.py

@@ -229,7 +229,7 @@ def search(
             List of search results with distance scores and payloads.
         """
         # Convert filter to Milvus expression
-        print(f"filter for milvus: {filter}")
+        logger.info(f"filter for milvus: {filter}")
         expr = self._dict_to_expr(filter) if filter else ""
 
         search_func_map = {