Skip to content

fix: add dummy_prefill guard for PD connection operations #3803

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 1, 2025

Conversation

FirwoodLin
Copy link
Contributor

Motivation

The current implementation attempts to call shelf_prefill_session on the PD (Prefill-Decode) connection pool regardless of whether dummy prefill mode is enabled. However, when dummy_prefill is set to True, there is no actual prefill session migration happening, and attempting to shelf a non-existent prefill session can lead to unnecessary operations or potential errors.

File "/nvme2/share/root/src/LMDeploy/lmdeploy/serve/proxy/proxy.py", line 643, in chat_completions_v1
    node_manager.pd_connection_pool.shelf_prefill_session((p_url, d_url), prefill_info['id'])
KeyError: 'id'

Modification

Added conditional guards around node_manager.pd_connection_pool.shelf_prefill_session() calls in two API endpoints:

  • /v1/chat/completions endpoint [chat_completions_v1]:
    Added if not node_manager.dummy_prefill: guard before calling shelf_prefill_session

  • /v1/completions endpoint [completions_v1]: Added if not node_manager.dummy_prefill: guard before calling shelf_prefill_session

Checklist

  1. Pre-commit or other linting tools are used to fix the potential lint issues.
  2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
  3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
  4. The documentation has been modified accordingly, like docstring or example tutorials.

@windreamer windreamer requested a review from grimoire August 1, 2025 02:32
Copy link
Collaborator

@grimoire grimoire left a comment

Choose a reason for hiding this comment

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

LGTM

@lvhan028 lvhan028 merged commit 57f1fd3 into InternLM:main Aug 1, 2025
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants