Skip to content

[MP][Feat] Support dedicated thread pool for MP callbacks#2763

Merged
KuntaiDu merged 2 commits intoLMCache:devfrom
ApostaC:local-dev/dedicate-worker-pool
Mar 15, 2026
Merged

[MP][Feat] Support dedicated thread pool for MP callbacks#2763
KuntaiDu merged 2 commits intoLMCache:devfrom
ApostaC:local-dev/dedicate-worker-pool

Conversation

@ApostaC
Copy link
Copy Markdown
Contributor

@ApostaC ApostaC commented Mar 13, 2026

What this PR does / why we need it:

Adds MessageQueueServer.add_dedicated_thread_pool(request_types, max_workers) so that specific blocking request types can run on their own ThreadPoolExecutor instead of sharing the single default pool. This prevents slow/blocking handlers (e.g., STORE, RETRIEVE) from starving other request types when the default pool is exhausted.

Usage (after handler registration, before server.start()):

server.add_dedicated_thread_pool([RequestType.STORE, RequestType.RETRIEVE], max_workers=4)

Multiple dedicated pools are supported. The method validates that all specified request types are registered BlockingRequestHandlers before creating the pool (two-pass for atomicity).

Special notes for your reviewers:

Minimal changes to existing code — only added one field to __init__, one guard in close(), and pool shutdown logic. The new method reassigns the executor attribute on already-registered BlockingRequestHandler instances. No existing methods were modified.

If applicable:

  • this PR contains user facing changes - docs added
  • this PR contains unit tests

@ApostaC
[add] dedicated worker pool support for MP mode
5dfb249 
Signed-off-by: ApostaC <yihua98@uchicago.edu>
@gemini-code-assist
Copy link
Copy Markdown
Contributor

gemini-code-assist Bot commented Mar 13, 2026

Summary of Changes

Hello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the MessageQueueServer by introducing the capability to assign dedicated thread pools to specific blocking request types. This architectural improvement addresses potential performance bottlenecks where long-running operations could exhaust the shared default thread pool, leading to service degradation for other request types. By allowing handlers to operate within their own isolated pools, the system gains improved responsiveness and resilience, ensuring that critical operations are not unduly impacted by less time-sensitive tasks.

Highlights

  • Dedicated Thread Pools: Introduced a new method MessageQueueServer.add_dedicated_thread_pool to allow specific blocking request types to run on their own ThreadPoolExecutor instances.
  • Concurrency Management: This feature prevents slow or blocking handlers (e.g., STORE, RETRIEVE) from monopolizing the default thread pool and starving other request types.
  • Validation and Safety: The new method includes validation to ensure that only registered BlockingRequestHandler instances can be assigned to dedicated pools, and supports multiple dedicated pools.
  • Resource Cleanup: Added proper shutdown logic for both the default and any dedicated thread pools when the MessageQueueServer is closed.
  • Comprehensive Testing: New unit tests were added to verify the functionality of dedicated thread pools, including correct assignment, error handling for invalid handler types or unregistered requests, and the ability to manage multiple pools.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • lmcache/v1/multiprocess/mq.py
    • Added dedicated_pools attribute to MessageQueueServer to store ThreadPoolExecutor instances.
    • Implemented add_dedicated_thread_pool method for assigning specific request types to new thread pools.
    • Modified the close method to ensure all dedicated thread pools are properly shut down.
  • tests/v1/multiprocess/test_mq.py
    • Imported BlockingRequestHandler and add_handler_helper for testing purposes.
    • Added test_add_dedicated_thread_pool to verify executor reassignment.
    • Included test_dedicated_thread_pool_error_on_sync_handler to check type validation.
    • Added test_dedicated_thread_pool_error_on_unregistered to test handler registration validation.
    • Implemented test_multiple_dedicated_thread_pools to confirm support for multiple dedicated pools.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist[bot]
gemini-code-assist Bot reviewed Mar 13, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a valuable feature for creating dedicated thread pools for specific request types, which will help prevent resource starvation from slow-running handlers. The implementation is clean and the two-pass validation in add_dedicated_thread_pool is a good approach for ensuring atomicity. The accompanying tests are thorough, covering success paths, error conditions, and the use of multiple dedicated pools. I have one suggestion to handle a minor edge case to prevent unnecessary resource allocation.

Comment thread lmcache/v1/multiprocess/mq.py
@ApostaC
fix gemini comments
af7d86d 
Signed-off-by: ApostaC <yihua98@uchicago.edu>
KuntaiDu
KuntaiDu approved these changes Mar 13, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@KuntaiDu KuntaiDu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@KuntaiDu KuntaiDu enabled auto-merge (squash) March 13, 2026 22:44
@github-actions github-actions Bot added the full Run comprehensive tests on this PR label Mar 13, 2026
@KuntaiDu KuntaiDu merged commit b42c3e4 into LMCache:dev Mar 15, 2026
28 of 29 checks passed
hyunyul-XCENA pushed a commit to xcena-dev/LMCache that referenced this pull request Mar 20, 2026
@ApostaC @hyunyul-XCENA
[MP][Feat] Support dedicated thread pool for MP callbacks (LMCache#2763)
705f4c6 
* [add] dedicated worker pool support for MP mode

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* fix gemini comments

Signed-off-by: ApostaC <yihua98@uchicago.edu>

---------

Signed-off-by: ApostaC <yihua98@uchicago.edu>
realAaronWu pushed a commit to realAaronWu/LMCache that referenced this pull request Mar 20, 2026
@ApostaC
[MP][Feat] Support dedicated thread pool for MP callbacks (LMCache#2763)
534c17d 
* [add] dedicated worker pool support for MP mode

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* fix gemini comments

Signed-off-by: ApostaC <yihua98@uchicago.edu>

---------

Signed-off-by: ApostaC <yihua98@uchicago.edu>
Signed-off-by: Aaron Wu <aaron.wu@dell.com>
deng451e pushed a commit to deng451e/LMCache that referenced this pull request Mar 25, 2026
@ApostaC
[MP][Feat] Support dedicated thread pool for MP callbacks (LMCache#2763)
5505d3c 
* [add] dedicated worker pool support for MP mode

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* fix gemini comments

Signed-off-by: ApostaC <yihua98@uchicago.edu>

---------

Signed-off-by: ApostaC <yihua98@uchicago.edu>
deng451e pushed a commit to deng451e/LMCache that referenced this pull request Mar 27, 2026
@ApostaC @deng451e
[MP][Feat] Support dedicated thread pool for MP callbacks (LMCache#2763)
5c39786 
* [add] dedicated worker pool support for MP mode

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* fix gemini comments

Signed-off-by: ApostaC <yihua98@uchicago.edu>

---------

Signed-off-by: ApostaC <yihua98@uchicago.edu>
jooho-XCENA pushed a commit to xcena-dev/LMCache that referenced this pull request Apr 2, 2026
@ApostaC @jooho-XCENA
[MP][Feat] Support dedicated thread pool for MP callbacks (LMCache#2763)
96cceb5 
* [add] dedicated worker pool support for MP mode

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* fix gemini comments

Signed-off-by: ApostaC <yihua98@uchicago.edu>

---------

Signed-off-by: ApostaC <yihua98@uchicago.edu>
jooho-XCENA pushed a commit to xcena-dev/LMCache that referenced this pull request Apr 2, 2026
@ApostaC @jooho-XCENA
[MP][Feat] Support dedicated thread pool for MP callbacks (LMCache#2763)
09c52ec 
* [add] dedicated worker pool support for MP mode

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* fix gemini comments

Signed-off-by: ApostaC <yihua98@uchicago.edu>

---------

Signed-off-by: ApostaC <yihua98@uchicago.edu>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@KuntaiDu KuntaiDu KuntaiDu approved these changes

@Oasis-Git Oasis-Git Oasis-Git approved these changes

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

full Run comprehensive tests on this PR

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ApostaC @KuntaiDu @Oasis-Git