scheduler feat: implementation of redis queue and new api search functions of mixture and fine mode (#462)

* debug an error function name

* feat: Add DynamicCache compatibility for different transformers versions

- Fix build_kv_cache method in hf.py to handle both old and new DynamicCache structures
  - Support new 'layers' attribute with key_cache/value_cache or keys/values
  - Maintain backward compatibility with direct key_cache/value_cache attributes
  - Add comprehensive error handling and logging for unsupported structures

- Update move_dynamic_cache_htod function in kv.py for cross-version compatibility
  - Handle layers-based structure in newer transformers versions
  - Support alternative attribute names (keys/values vs key_cache/value_cache)
  - Preserve original functionality for older transformers versions

- Add comprehensive tests for DynamicCache compatibility
  - Test activation memory update with mock DynamicCache layers
  - Verify layers attribute access across different transformers versions
  - Fix scheduler logger mock to include memory_manager attribute
This resolves AttributeError issues when using different versions of the
transformers library and ensures robust handling of DynamicCache objects.

debug

* feat: implement APIAnalyzerForScheduler for memory operations

- Add APIAnalyzerForScheduler class with search/add operations
- Support requests and http.client with connection reuse
- Include comprehensive error handling and dynamic configuration
- Add English test suite with real-world conversation scenarios

* feat: Add search_ws API endpoint and enhance API analyzer functionality

- Add search_ws endpoint in server_router.py for scheduler-enabled search
- Fix missing imports: time module, SearchRequest class, and get_mos_product_instance function
- Implement search_ws method in api_analyzer.py with HTTP client support
- Add _search_ws_with_requests and _search_ws_with_http_client private methods
- Include search_ws usage example in demonstration code
- Enhance scheduler and dispatcher capabilities for improved memory management
- Expand test coverage to ensure functionality stability

This update primarily strengthens the memory scheduling system's search capabilities,
providing users with more flexible API interface options.

* fix: resolve test failures and warnings in test suite

- Fix Pydantic serialization warning in test_memos_chen_tang_hello_world
  * Add warnings filter to suppress UserWarning from Pydantic serialization

- Fix KeyError: 'past_key_values' in test_build_kv_cache_and_generation
  * Update mock configuration to properly return forward_output with past_key_values
  * Add DynamicCache version compatibility handling in test mocks
  * Support both old and new transformers versions with layers/key_cache attributes
  * Improve assertion logic to check all model calls for required parameters

- Update base_scheduler.py to use centralized DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE constant
  * Add import for DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE from general_schemas
  * Replace hardcoded value 100 with configurable constant (1000)

All tests now pass successfully with proper version compatibility handling.

* feat: add a test_robustness execution to test thread pool execution

* feat: optimize scheduler configuration and API search functionality

- Add DEFAULT_TOP_K and DEFAULT_CONTEXT_WINDOW_SIZE global constants in general_schemas.py
- Update base_scheduler.py to use global default values instead of hardcoded numbers
- Fix SchedulerConfigFactory initialization issue by using keyword argument expansion
- Resolve UnboundLocalError variable conflict in search_memories_ws function
- Fix indentation and parameter issues in OptimizedScheduler search_for_api method
- Improve code standardization and maintainability

* feat: Add Redis auto-initialization with fallback strategies

- Add auto_initialize_redis() with config/env/local fallback
- Move Redis logic from dispatcher_monitor to redis_service
- Update base_scheduler to use auto initialization
- Add proper resource cleanup and error handling

* feat: add database connection management to ORM module

- Add MySQL engine loading from environment variables in BaseDBManager
- Add Redis connection loading from environment variables in BaseDBManager
- Enhance database configuration validation and error handling
- Complete database adapter infrastructure for ORM module
- Provide unified database connection management interface

This update provides comprehensive database connection management capabilities
for the mem_scheduler module, supporting dynamic MySQL and Redis configuration
loading from environment variables, establishing reliable data persistence
foundation for scheduling services and API services.

* remove part of test

* feat: add Redis-based ORM with multiprocess synchronization

- Add RedisDBManager and RedisLockableORM classes
- Implement atomic locking mechanism for concurrent access
- Add merge functionality for different object types
- Include comprehensive test suite and examples
- Fix Redis key type conflicts in lock operations

* fix: resolve scheduler module import and Redis integration issues

* revise naive memcube creation in server router

* remove long-time tests in test_scheduler

* remove redis test which needs .env

* refactor all codes about mixture search with scheduler

* fix: resolve Redis API synchronization issues and implement search API with reranker

- Fix running_entries to running_task_ids migration across codebase
- Update sync_search_data method to properly handle TaskRunningStatus
- Correct variable naming and logic in API synchronization flow
- Implement search API endpoint with reranker functionality
- Update test files to reflect new running_task_ids convention
- Ensure proper Redis state management for concurrent tasks

* remove a test for api module

* revise to pass the test suite

* address some bugs to make mix_search normally running

* modify codes according to evaluation logs

* feat: Optimize mixture search and enhance API client

* feat: Add conversation_turn tracking for session-based memory search

    - Add conversation_turn field to APIMemoryHistoryEntryItem schema with default value 0
    - Implement session counter in OptimizedScheduler to track turn count per session_id
    - Update sync_search_data method to accept and store conversation_turn parameter
    - Maintain session history with LRU eviction (max 5 sessions)
    - Rename conversation_id to session_id for consistency with request object
    - Enable direct access to session_id from search requests

    This feature allows tracking conversation turns within the same session,
    providing better context for memory retrieval and search history management.

* adress time bug in monitor

* revise simple tree

* add mode to evaluation client; rewrite print to logger.info in db files

* feat: 1. add redis queue for scheduler 2. finish the code related to mix search and fine search

* debug the working memory code

* addressed a range of bugs to make scheduler running correctly

* remove test_dispatch_parallel test

* print change to logger.info

---------

Co-authored-by: CaralHsi <[email protected]>

Author: tangg555

examples/mem_scheduler/api_w_scheduler.py

@@ -0,0 +1,62 @@
+from memos.api.routers.server_router import mem_scheduler
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+
+
+# Debug: Print scheduler configuration
+print("=== Scheduler Configuration Debug ===")
+print(f"Scheduler type: {type(mem_scheduler).__name__}")
+print(f"Config: {mem_scheduler.config}")
+print(f"use_redis_queue: {mem_scheduler.use_redis_queue}")
+print(f"Queue type: {type(mem_scheduler.memos_message_queue).__name__}")
+print(f"Queue maxsize: {getattr(mem_scheduler.memos_message_queue, 'maxsize', 'N/A')}")
+
+# Check if Redis queue is connected
+if hasattr(mem_scheduler.memos_message_queue, "_is_connected"):
+    print(f"Redis connected: {mem_scheduler.memos_message_queue._is_connected}")
+if hasattr(mem_scheduler.memos_message_queue, "_redis_conn"):
+    print(f"Redis connection: {mem_scheduler.memos_message_queue._redis_conn}")
+print("=====================================\n")
+
+queue = mem_scheduler.memos_message_queue
+queue.clear()
+
+
+# 1. Define a handler function
+def my_test_handler(messages: list[ScheduleMessageItem]):
+    print(f"My test handler received {len(messages)} messages:")
+    for msg in messages:
+        print(f" my_test_handler - {msg.item_id}: {msg.content}")
+        print(
+            f"{queue._redis_conn.xinfo_groups(queue.stream_name)} qsize: {queue.qsize()} messages:{messages}"
+        )
+
+
+# 2. Register the handler
+TEST_HANDLER_LABEL = "test_handler"
+mem_scheduler.register_handlers({TEST_HANDLER_LABEL: my_test_handler})
+
+# 3. Create messages
+messages_to_send = [
+    ScheduleMessageItem(
+        item_id=f"test_item_{i}",
+        user_id="test_user",
+        mem_cube_id="test_mem_cube",
+        label=TEST_HANDLER_LABEL,
+        content=f"This is test message {i}",
+    )
+    for i in range(5)
+]
+
+# 5. Submit messages
+for mes in messages_to_send:
+    print(f"Submitting message {mes.item_id} to the scheduler...")
+    mem_scheduler.submit_messages([mes])
+
+# 6. Wait for messages to be processed (limited to 100 checks)
+print("Waiting for messages to be consumed (max 100 checks)...")
+mem_scheduler.mem_scheduler_wait()
+
+
+# 7. Stop the scheduler
+print("Stopping the scheduler...")
+mem_scheduler.stop()

examples/mem_scheduler/memos_w_optimized_scheduler.py

@@ -70,13 +70,48 @@ def init_task():
     return conversations, questions
 
 
+def show_web_logs(mem_scheduler: GeneralScheduler):
+    """Display all web log entries from the scheduler's log queue.
+
+    Args:
+        mem_scheduler: The scheduler instance containing web logs to display
+    """
+    if mem_scheduler._web_log_message_queue.empty():
+        print("Web log queue is currently empty.")
+        return
+
+    print("\n" + "=" * 50 + " WEB LOGS " + "=" * 50)
+
+    # Create a temporary queue to preserve the original queue contents
+    temp_queue = Queue()
+    log_count = 0
+
+    while not mem_scheduler._web_log_message_queue.empty():
+        log_item: ScheduleLogForWebItem = mem_scheduler._web_log_message_queue.get()
+        temp_queue.put(log_item)
+        log_count += 1
+
+        # Print log entry details
+        print(f"\nLog Entry #{log_count}:")
+        print(f'- "{log_item.label}" log: {log_item}')
+
+        print("-" * 50)
+
+    # Restore items back to the original queue
+    while not temp_queue.empty():
+        mem_scheduler._web_log_message_queue.put(temp_queue.get())
+
+    print(f"\nTotal {log_count} web log entries displayed.")
+    print("=" * 110 + "\n")
+
+
 def run_with_scheduler_init():
     print("==== run_with_automatic_scheduler_init ====")
     conversations, questions = init_task()
 
     # set configs
     mos_config = MOSConfig.from_yaml_file(
-        f"{BASE_DIR}/examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml"
+        f"{BASE_DIR}/examples/data/config/mem_scheduler/memos_config_w_optimized_scheduler.yaml"
     )
 
     mem_cube_config = GeneralMemCubeConfig.from_yaml_file(
@@ -118,6 +153,7 @@ def run_with_scheduler_init():
     )
 
     mos.add(conversations, user_id=user_id, mem_cube_id=mem_cube_id)
+    mos.mem_scheduler.current_mem_cube = mem_cube
 
     for item in questions:
         print("===== Chat Start =====")
@@ -131,40 +167,5 @@ def run_with_scheduler_init():
     mos.mem_scheduler.stop()
 
 
-def show_web_logs(mem_scheduler: GeneralScheduler):
-    """Display all web log entries from the scheduler's log queue.
-
-    Args:
-        mem_scheduler: The scheduler instance containing web logs to display
-    """
-    if mem_scheduler._web_log_message_queue.empty():
-        print("Web log queue is currently empty.")
-        return
-
-    print("\n" + "=" * 50 + " WEB LOGS " + "=" * 50)
-
-    # Create a temporary queue to preserve the original queue contents
-    temp_queue = Queue()
-    log_count = 0
-
-    while not mem_scheduler._web_log_message_queue.empty():
-        log_item: ScheduleLogForWebItem = mem_scheduler._web_log_message_queue.get()
-        temp_queue.put(log_item)
-        log_count += 1
-
-        # Print log entry details
-        print(f"\nLog Entry #{log_count}:")
-        print(f'- "{log_item.label}" log: {log_item}')
-
-        print("-" * 50)
-
-    # Restore items back to the original queue
-    while not temp_queue.empty():
-        mem_scheduler._web_log_message_queue.put(temp_queue.get())
-
-    print(f"\nTotal {log_count} web log entries displayed.")
-    print("=" * 110 + "\n")
-
-
 if __name__ == "__main__":
     run_with_scheduler_init()