refactor: scheduler #957
Open
refactor: scheduler #957
+3,066
−1,707
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
… fast search when api_module is None
…iever - Split SchedulerRetriever into SchedulerSearchService and MemoryPostProcessor - SchedulerSearchService: Unified search interface delegating to Searcher - MemoryPostProcessor: Handles all post-retrieval processing (enhance, filter, rerank) - Fix: Use memory_type='All' to prevent 2*top_k bug in search results - Add robust search_method conversion with logging for unknown methods - Update all scheduler call sites (base, general, optimized, eval) - Add comprehensive unit tests for both new services Breaking changes: None (backward compatible) Fixes: Search result count bug (was returning up to 2*top_k)
Problem: - enhance_memories_with_query and recall_for_missing_memories were in MemoryPostProcessor - This caused SearchHandler to depend on Scheduler (cross-layer coupling) - These are general search enhancement features, not scheduler-specific Solution: - Move both methods from MemoryPostProcessor to AdvancedSearcher - Update SingleCubeView to use self.searcher instead of self.mem_scheduler.post_processor - Simplify MemoryPostProcessor to only handle filtering and reranking Benefits: - Better separation of concerns (search enhancement belongs to Searcher) - Removes cross-layer dependency (API layer no longer depends on Scheduler) - Enables reuse (both Scheduler and SearchHandler share same implementation) - Improves testability (AdvancedSearcher can be tested independently) Changes: - advanced_searcher.py: +291 lines (added enhancement methods) - post_processor.py: -320 lines (removed enhancement, kept filter/rerank) - single_cube.py: 6 changes (updated call sites) - test_post_processor.py: -57 lines (removed obsolete tests) Type safety: Verified via 'AdvancedSearcher as Searcher' alias in tree.py
Fixes 3 critical issues identified in code review: 1. Configuration behavior change (HIGH RISK) Problem: enhance_memories_with_query called without batch_size/retries Impact: Behavior inconsistent with historical settings Fix: Pass DEFAULT_SCHEDULER_RETRIEVER_BATCH_SIZE and RETRIES explicitly 2. Module dependency inversion (ARCHITECTURE REGRESSION) Problem: AdvancedSearcher depended on memos.mem_scheduler.utils Impact: Violated decoupling goal (memories should not depend on scheduler) Fix: Move extract_json_obj/extract_list_items_in_answer to memos.utils Changes: - src/memos/utils.py: +163 lines (add utility functions) - misc_utils.py: -163 lines, +3 lines (re-export for backward compat) - Update imports in 6 files 3. Test coverage migration (NOTED) Note: _split_batches tests removed but not migrated Action: Defer to future test enhancement Benefits: - Maintains historical batch/retry behavior - Restores proper layering (no memories→scheduler dependency) - Backward compatible (misc_utils re-exports from memos.utils) Verification: ✓ No scheduler imports in AdvancedSearcher ✓ Config params passed to enhance_memories_with_query ✓ All existing imports work via re-export
Critical fix for misc_utils.py: - Removed empty 'def extract_json_obj(text: str):' declaration - This was leftover from sed deletion and caused IndentationError - Functions are properly re-exported from memos.utils Verification: python3 -m py_compile passes
- Fix test_search_service.py: call_count assertion (2 -> 1) SearchService now uses memory_type='All' for single search call - Fix test_post_processor.py: remove incorrect 'with Mock()' usage - Fix ruff E712: use truthiness check instead of '== True' - Apply ruff formatting to all modified files All 23 tests now pass successfully.
Refactor scheduler retriever
- Extract activation memory management logic from `BaseScheduler` to new `ActivationMemoryManager` class - Move background task schedule monitoring loop from `BaseScheduler` to `TaskScheduleMonitor` - Update `BaseScheduler` to delegate activation memory operations to `ActivationMemoryManager` - Update `BaseScheduler` to use `TaskScheduleMonitor` for metrics monitoring - Update tests to reflect changes in scheduler component naming (retriever -> post_processor/search_service)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Summary: (summary)
Key Refactoring Changes:
Handlers Decomposition & Decoupling (Handlers 拆分与解耦) :
Unified Monitoring Logic (监控逻辑统一) :
Unified Retrieval Logic (召回逻辑统一) :
Redis Module Decomposition (Redis 模块拆分) :
Unified Error Handling (错误处理与重试机制统一) :
Fix: #(issue)
Docs Issue/PR: (docs-issue-or-pr-link)
Reviewer: @(reviewer)
Checklist: