Commit 9341861
Feat/redis scheduler (#526)
* 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
* adjucted the core code related to fine and mixture apis
* feat: create task queue to wrap local queue and redis queue. queue now split FIFO to multi queue from different users. addressed a range of bugs
* fix bugs: debug bugs about internet trigger
* debug get searcher mode
* feat: add manual internet
* Fix: fix code format
* feat: add strategy for fine search
* debug redis queue
* debug redis queue
* fix bugs: completely addressed bugs about redis queue
* refactor: add searcher to handler_init; remove info log from task_queue
* refactor: modify analyzer
* refactor: revise locomo_eval to make it support llm other than gpt-4o-mini
* feat: develop advanced searcher with deep search
* feat: finish a complete version of deep search
* refactor: refactor deep search feature, now only allowing one-round deep search
* feat: implement the feature of get_tasks_status, but completed tasks are not recorded yet; waiting to be developed
* debuging merged code; searching memories have bugs
* change logging level
* debug api evaluation
* fix bugs: change top to top_k
* change log
* feat(scheduler): Implement comprehensive observability and fix critical bugs
This commit introduces a robust observability system for the scheduler and resolves several critical bugs identified during code review and testing.
Key Improvements:
- **Task Status Tracking**: Implemented `TaskStatusTracker` using Redis to provide persistent, per-task lifecycle tracking (`waiting`, `in_progress`, `completed`, `failed`).
- **Prometheus Metrics**: Added a new metrics system to expose key performance indicators (QPS, latency, queue length, failure/completion rates) for monitoring.
- **API Refactoring**: Refactored `/scheduler/status` and `/scheduler/wait` APIs to use the new reliable `TaskStatusTracker`, ensuring accurate state reporting.
Bug Fixes:
- **Initialization**: Corrected the `SchedulerDispatcher` initialization order to prevent `NoneType` errors in tests and at runtime.
- **CPU Usage**: Fixed a busy-wait loop in the metrics monitor thread that caused 100% CPU usage when idle.
- **Exception Handling**: Refined API handlers to correctly propagate HTTP error codes (e.g., 404) instead of masking them as 500 errors.
- **Dependencies**: Added missing dependencies (`prometheus-client`) and updated test mocks to ensure the test suite can run correctly.
- **Legacy Code**: Removed the old, buggy `mem_scheduler_wait` method.
All 394 unit tests now pass, and a functional test of the new features has been successfully verified.
* fix(ci): Resolve top-level redis import error in TaskStatusTracker
* feat(scheduler): Implement conditional cloud status updates
Adds functionality to send task status updates (success/failure) to RabbitMQ via , specifically for the cloud service platform. This includes:
- Adding a field to .
- Passing to .
- Modifying to conditionally send with or (along with and ) based on the environment variable.
* fix(deps): Promote prometheus-client to core dependency
Moves prometheus-client from an optional group to the main project dependencies. This ensures it is always installed in all environments, including the CI/CD deployment pipeline, to resolve the recurring 'ModuleNotFoundError'.
* fix(ci): Resolve ruff linting and style errors
Addresses all linting errors reported by ruff, including undefined names for 'os' and 'timezone' by correcting import statements in 'dispatcher.py' and 'status_tracker.py'. Also resolves various code style violations (line breaks, comment punctuation) to align with ruff standards. This ensures the code adheres to project standards and passes the CI quality checks.
* fix(ci): Reformat code to comply with ruff standards
* fix(docker): Add prometheus-client to Docker requirements.txt
Adds prometheus-client to docker/requirements.txt to ensure it is installed in the Docker build environment, resolving deployment failures related to missing dependencies.
---------
Co-authored-by: chentang <[email protected]>
Co-authored-by: fridayL <[email protected]>
Co-authored-by: [email protected] <>1 parent a921e51 commit 9341861
File tree
37 files changed
+1828
-2113
lines changed- docker
- evaluation/scripts
- locomo
- utils
- examples/mem_scheduler
- src/memos
- api
- handlers
- routers
- mem_scheduler
- analyzer
- memory_manage_modules
- monitors
- schemas
- task_schedule_modules
- utils
- memories/textual
- tree_text_memory/retrieve
- multi_mem_cube
- templates
- types
- tests
- api
- mem_scheduler
37 files changed
+1828
-2113
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
158 | 158 | | |
159 | 159 | | |
160 | 160 | | |
| 161 | + | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
3 | 3 | | |
4 | 4 | | |
5 | 5 | | |
| 6 | + | |
6 | 7 | | |
7 | 8 | | |
8 | 9 | | |
| |||
47 | 48 | | |
48 | 49 | | |
49 | 50 | | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
| 56 | + | |
| 57 | + | |
| 58 | + | |
| 59 | + | |
| 60 | + | |
| 61 | + | |
| 62 | + | |
| 63 | + | |
| 64 | + | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
50 | 74 | | |
51 | 75 | | |
52 | 76 | | |
| |||
77 | 101 | | |
78 | 102 | | |
79 | 103 | | |
80 | | - | |
81 | | - | |
82 | | - | |
83 | | - | |
84 | | - | |
85 | | - | |
86 | | - | |
87 | | - | |
88 | | - | |
89 | | - | |
90 | | - | |
91 | | - | |
92 | | - | |
93 | | - | |
| 104 | + | |
| 105 | + | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
| 109 | + | |
| 110 | + | |
| 111 | + | |
| 112 | + | |
| 113 | + | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
| 117 | + | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
94 | 121 | | |
95 | 122 | | |
96 | 123 | | |
| |||
284 | 311 | | |
285 | 312 | | |
286 | 313 | | |
287 | | - | |
| 314 | + | |
288 | 315 | | |
289 | 316 | | |
290 | 317 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
189 | 189 | | |
190 | 190 | | |
191 | 191 | | |
192 | | - | |
193 | | - | |
194 | | - | |
| 192 | + | |
195 | 193 | | |
196 | 194 | | |
197 | 195 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
1 | 3 | | |
2 | 4 | | |
3 | 5 | | |
4 | 6 | | |
5 | | - | |
| 7 | + | |
6 | 8 | | |
7 | 9 | | |
8 | 10 | | |
| |||
26 | 28 | | |
27 | 29 | | |
28 | 30 | | |
29 | | - | |
| 31 | + | |
30 | 32 | | |
31 | | - | |
| 33 | + | |
32 | 34 | | |
33 | 35 | | |
34 | 36 | | |
35 | 37 | | |
| 38 | + | |
36 | 39 | | |
37 | 40 | | |
38 | 41 | | |
39 | | - | |
40 | | - | |
41 | | - | |
| 42 | + | |
42 | 43 | | |
43 | 44 | | |
44 | 45 | | |
45 | 46 | | |
46 | 47 | | |
47 | 48 | | |
48 | | - | |
| 49 | + | |
49 | 50 | | |
50 | 51 | | |
51 | 52 | | |
| |||
56 | 57 | | |
57 | 58 | | |
58 | 59 | | |
59 | | - | |
| 60 | + | |
| 61 | + | |
60 | 62 | | |
61 | 63 | | |
62 | 64 | | |
63 | 65 | | |
64 | 66 | | |
65 | | - | |
66 | | - | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
67 | 70 | | |
68 | 71 | | |
69 | 72 | | |
70 | | - | |
| 73 | + | |
| 74 | + | |
71 | 75 | | |
72 | | - | |
73 | | - | |
| 76 | + | |
74 | 77 | | |
75 | 78 | | |
76 | 79 | | |
77 | 80 | | |
78 | | - | |
79 | | - | |
80 | | - | |
| 81 | + | |
81 | 82 | | |
82 | 83 | | |
83 | 84 | | |
| |||
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
46 | 46 | | |
47 | 47 | | |
48 | 48 | | |
| 49 | + | |
49 | 50 | | |
50 | 51 | | |
51 | 52 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
9 | 9 | | |
10 | 10 | | |
11 | 11 | | |
12 | | - | |
| 12 | + | |
13 | 13 | | |
14 | 14 | | |
15 | 15 | | |
| |||
132 | 132 | | |
133 | 133 | | |
134 | 134 | | |
135 | | - | |
| 135 | + | |
136 | 136 | | |
137 | 137 | | |
138 | 138 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
129 | 129 | | |
130 | 130 | | |
131 | 131 | | |
| 132 | + | |
| 133 | + | |
| 134 | + | |
| 135 | + | |
| 136 | + | |
| 137 | + | |
| 138 | + | |
| 139 | + | |
| 140 | + | |
| 141 | + | |
| 142 | + | |
| 143 | + | |
| 144 | + | |
| 145 | + | |
| 146 | + | |
132 | 147 | | |
133 | 148 | | |
134 | 149 | | |
| |||
272 | 287 | | |
273 | 288 | | |
274 | 289 | | |
| 290 | + | |
| 291 | + | |
275 | 292 | | |
276 | 293 | | |
277 | 294 | | |
| |||
286 | 303 | | |
287 | 304 | | |
288 | 305 | | |
| 306 | + | |
289 | 307 | | |
290 | 308 | | |
291 | 309 | | |
| |||
335 | 353 | | |
336 | 354 | | |
337 | 355 | | |
| 356 | + | |
338 | 357 | | |
339 | 358 | | |
0 commit comments