Support updating weights at once by stopping all requests

[tianyuzhou95]

Motivation

Currently, when updating model weights, we have to wait for all requests to complete generation before we can start updating. This significantly reduces the throughput of batch inference especially in RL scenarios. Therefore, a method is needed to quickly stop all running/queued requests so that model weights can be updated immediately.

Modifications

1. Return result rather than silently ignore when abort request in waiting queue.
2. Extend the native API /abort_request to abort all requests.
3. Add a new boolean field abort_all_requests inside the request of updating weights to decide whether to abort all requests before updating weights.

[gemini-code-assist]

Hello @tianyuzhou95, I'm Gemini Code Assist¹! 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!

Summary of Changes

Hello team! Gemini here, providing a summary of this pull request to help everyone get up to speed quickly.

This PR introduces functionality to quickly stop all ongoing and queued requests within the SGLang runtime. The primary motivation, as described by the author, is to enable immediate model weight updates without having to wait for potentially long-running requests to complete, which is particularly beneficial for use cases like batch inference in RL scenarios where throughput is critical.

The changes involve adding a new native HTTP API endpoint /stop_all that can be called directly, and also integrating this stop mechanism into the existing weight update APIs (update_weights_from_disk, update_weights_from_distributed, update_weights_from_tensor) via a new boolean parameter stop_all_requests. When triggered, the system will modify the maximum number of new tokens for both waiting and running requests, effectively causing them to finish immediately or at their current state.

Highlights

• New API Endpoint: Adds a new /stop_all HTTP API endpoint to trigger the stopping of all requests.
• Weight Update Integration: Adds a stop_all_requests boolean field to the weight update request inputs, allowing requests to be stopped automatically before a weight update.
• Scheduler Logic: Implements the core logic in the scheduler to handle the StopAllReq, modifying the max_new_tokens for queued and running requests to halt generation.

Changelog

Click here to see the changelog

• python/sglang/srt/entrypoints/http_server.py
  • Added /stop_all API endpoint (GET/POST) that calls _global_state.tokenizer_manager.stop_all().
• python/sglang/srt/managers/io_struct.py
  • Added StopAllReq dataclass.
  • Added stop_all_requests: bool = False field to UpdateWeightFromDiskReqInput.
  • Added stop_all_requests: bool = False field to UpdateWeightsFromDistributedReqInput.
  • Added stop_all_requests: bool = False field to UpdateWeightsFromTensorReqInput.
• python/sglang/srt/managers/scheduler.py
  • Imported StopAllReq.
  • Added handler for StopAllReq in the scheduler's request mapping.
  • Implemented stop_all_request method to set max_new_tokens to 0 for waiting requests and to the current output length for running requests.
• python/sglang/srt/managers/tokenizer_manager.py
  • Imported StopAllReq.
  • Added stop_all() method that sends a StopAllReq to the scheduler.
  • Added checks in update_weights_from_disk, update_weights_from_distributed, and update_weights_from_tensor to call self.stop_all() if stop_all_requests is true.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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.

---

Requests flow fast,
But models need update,
Stop them all now.

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]

Code Review

This pull request introduces a valuable feature to pause/stop ongoing requests before model weights are updated, which is particularly useful for scenarios like reinforcement learning. The approach involves a new API endpoint /stop_all and flags in weight update requests to trigger this behavior. The mechanisms for marking requests to stop within the scheduler by adjusting max_new_tokens seem correct.

The main concern revolves around the synchronization: ensuring that requests are effectively stopped or paused in the scheduler before the weight update operation commences. The current implementation sends a stop signal asynchronously, which might lead to race conditions where weights are updated while requests are still being processed with old weights, or caches are flushed prematurely. I've provided a detailed comment and suggestions on this critical aspect.

Minor points include a variable naming choice, which is noted in the summary but not as an individual comment due to review settings.

Summary of Findings

• Synchronization of Request Stopping and Weight Updates: Critical: The current mechanism for stopping requests (stop_all()) sends an asynchronous message to the scheduler. There's no guarantee that requests are stopped before the weight update process begins, potentially leading to race conditions or failed cache flushes. This needs to be made synchronous or ensure quiescence before updates.
• Variable Naming: Medium: In scheduler.py, the variable sum shadows the built-in function, potentially causing confusion.
• Clarity of /stop_all API behavior: Medium: The /stop_all API endpoint's current fire-and-forget nature might not align with user expectations of immediate, confirmed stopping. Its behavior should be clarified or made more robust if tokenizer_manager.stop_all becomes async.
• Indentation: Low: In tokenizer_manager.py, the line req = StopAllReq() has inconsistent indentation (1 space instead of 4). (Not commented due to severity settings)

Merge Readiness

This pull request introduces an important capability. However, due to the critical synchronization issue identified between stopping requests and updating model weights, I recommend that these changes not be merged until this issue is thoroughly addressed. Addressing this will ensure the stability and correctness of the dynamic weight update feature. The other medium severity suggestions should also be considered for improving clarity and robustness. I am not authorized to approve pull requests, so please ensure further review and approval by authorized maintainers after addressing the feedback.

[tianyuzhou95]

@zhuzilin seems this PR also meets your case? cc @zhaochenyang20

related pr: #6855

[zhaochenyang20]

@tianyuzhou95 please resolve the conflcits. I will merge #6855 and #6698

[tianyuzhou95]

@tianyuzhou95 please resolve the conflcits. I will merge #6855 and #6698

@zhaochenyang20 Thanks, already resolved :)

[zhyncs]

please rebase

[tianyuzhou95]

please rebase

Thanks, already rebased.

[zhaochenyang20]

Add a new boolean field abort_all_requests inside the request of updating weights to decide whether to abort all requests before updating weights.

That means if this is turn off, the update will wait all the requests to finish, then update?

[tianyuzhou95]

Add a new boolean field abort_all_requests inside the request of updating weights to decide whether to abort all requests before updating weights.

That means if this is turn off, the update will wait all the requests to finish, then update?

Yes, please refer https://github.com/sgl-project/sglang/blob/v0.4.8/python/sglang/srt/managers/tokenizer_manager.py#L889, the update will be blocked by the write lock.

[hebiao064]

will ask @zhyncs and @zhaochenyang20 to merge this PR once CI passed, thanks a lot!

[hebiao064]

pls fix lint by pre-commit run --all-files

[tianyuzhou95]

pre-commit run --all-files

Thanks, already fix linting. Besides, also add an unit test for updating weights with abort_all_requests=True.

[zhyncs]

Hi @tianyuzhou95 please rebase cc @zhuzilin

[tianyuzhou95]

Hi @tianyuzhou95 please rebase cc @zhuzilin

Thanks for pointing out, already rebased :)