[WIP][v1] Support for returning a value when using wait_for_save

[flesher0813]

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Current connectors do not specify whether wait_for_save returns a value.
This PR adds an explicit return value so that workers can inform the scheduler whether the request’s blocks were dumped successfully.

Test Plan

Test with test_multi_connector.py

Test Result

TODO

(Optional) Documentation Update

[gemini-code-assist]

Code Review

This pull request introduces support for a return value from wait_for_save to notify the scheduler about successfully dumped blocks. The changes are logical, but I've identified a critical bug in the scheduler where the new logic is outside the loop it depends on, which would lead to incorrect behavior. I've also suggested adding type hints in a few places to improve code clarity and maintainability.

[gemini-code-assist]

This block of code appears to be outside the for request in self.running: loop, but it uses the loop variables request and req_id. This will cause the logic to only apply to the last request processed in the loop, which is likely a bug. This block should be moved inside the for loop to ensure it's executed for every running request.

Additionally, the indentation on the second line of this block is inconsistent.

Suggested change

	if model_runner_output.finished_dumping is not None:
	request.succeed_dumped_blocks.extend(model_runner_output.finished_dumping.get(req_id, []))
	if model_runner_output.finished_dumping is not None:
	request.succeed_dumped_blocks.extend(model_runner_output.finished_dumping.get(req_id, []))

[gemini-code-assist]

The new attribute succeed_dumped_blocks is initialized without a type hint. For consistency and type safety, please add one. Based on its usage elsewhere, the type should be list[str].

Suggested change

	self.succeed_dumped_blocks = []
	self.succeed_dumped_blocks: list[str] = []

[gemini-code-assist]

The return type hint for maybe_wait_for_kv_save was removed when the function was changed to return a value. To maintain code clarity and type safety, please add the appropriate return type hint. Based on the type of finished_dumping in ModelRunnerOutput, it should be Optional[dict[str, list[str]]].

Suggested change

	def maybe_wait_for_kv_save():
	def maybe_wait_for_kv_save() -> Optional[dict[str, list[str]]]:

[github-actions]

👋 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 fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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 ready label to the PR or enable auto-merge.

🚀