v1: Pass KVConnectorOutput to scheduler-side#22157
v1: Pass KVConnectorOutput to scheduler-side#22157vllm-bot merged 1 commit intovllm-project:mainfrom
Conversation
orozery
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
|
cc @njhill |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new scheduler-side connector API, update_connector_output, to allow the connector to process outputs from worker-side connectors. The changes are well-contained and logically sound. I have one suggestion to improve the robustness of the new API by clarifying its contract in the docstring to prevent potential issues in future implementations.
There was a problem hiding this comment.
It's good practice to clarify the contract of this method regarding modifications to the connector_output object. The _update_from_kv_xfer_finished method in scheduler.py iterates over connector_output.finished_recving and connector_output.finished_sending after calling this method. If an implementation of update_connector_output were to modify these fields (e.g., clear them), the scheduler would miss these finished transfers, leading to bugs like requests getting stuck or memory leaks.
A similar method, build_connector_meta, has a note in its docstring: This function should NOT modify fields in the scheduler_output.
I recommend adding a similar note here to prevent potential issues with future connector implementations.
| Update KVConnector state from worker-side connectors output. | |
| Update KVConnector state from worker-side connectors output. | |
| NOTE: This function should not modify the `connector_output` object, | |
| as its fields may be used by the scheduler after this call. | |
There was a problem hiding this comment.
I'm actually thinking on allowing this...
but waiting for human feedback on that
There was a problem hiding this comment.
I actually agree with the bot here. See my other comment below for more details.
orozery
force-pushed
the
update-kv-connector-output
branch
from
1e62275 to
7a9122e
Compare
This commit adds a new scheduler-side connector v1 API function to notify the scheduler-side connector on output from the worker-side connectors. In particular, it allows the scheduler-side connector to track the finished requests. Signed-off-by: Or Ozeri <oro@il.ibm.com>
orozery
force-pushed
the
update-kv-connector-output
branch
from
7a9122e to
c92e866
Compare
Good point. |
|
|
||
| assert self.connector is not None | ||
| self.connector.update_connector_output(kv_connector_output) | ||
|
|
There was a problem hiding this comment.
I think update_connector_output should be invoked directly from update_from_output rather than from _update_from_kv_xfer_finished.
The issue is that _update_from_kv_xfer_finished only processes finished requests, while other (future) parts of kv_connector_output may be handled elsewhere in update_from_output — for example, invalid_block_ids in the KV load failure recovery PR (#19330).
Invoking update_connector_output independently would keep responsibilities clearer and help avoid unexpected side effects. For the same reason, I think kv_connector_output should remain unmodified by update_connector_output.
There was a problem hiding this comment.
@sdavidbd I agree.
I was trying to recall why in the first place I wanted to be able to modify the connector output.
I think it was because I was targeting a case where my workers will report block-wise results, and I wanted my scheduler-side connector to aggregate and produce request-based results to the scheduler.
But I think we all lean towards making the scheduler itself block-based. Then we don't need it.
So anyhow, I'm also good with changing this function to not allow to update the output (is there anything better than adding a comment?)
Signed-off-by: Or Ozeri <oro@il.ibm.com> Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Or Ozeri <oro@il.ibm.com> Signed-off-by: Diego-Castan <diego.castan@ibm.com>
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Signed-off-by: Or Ozeri <oro@il.ibm.com> Signed-off-by: Xiao Yu <xiao.yu@amd.com>
Signed-off-by: Or Ozeri <oro@il.ibm.com>
4 participants
This PR adds a new scheduler-side connector v1 API function to notify the scheduler-side connector on output from the worker-side connectors. In particular, it allows the scheduler-side connector to track the finished requests.
Part of the work described in RFC #19854