diff --git a/docs/en/advance/turbomind_priority_scheduling.md b/docs/en/advance/turbomind_priority_scheduling.md new file mode 100644 index 0000000000..11e7918cdc --- /dev/null +++ b/docs/en/advance/turbomind_priority_scheduling.md @@ -0,0 +1,125 @@ +# TurboMind Priority Scheduling + +TurboMind supports request-level priority scheduling. This is useful when one service handles different classes of traffic at the same time, such as interactive requests, background batch jobs, regular users, and high-priority users. After priority scheduling is enabled, requests with higher priority are admitted to inference first, while requests that have already started are kept whenever possible to reduce extra KV cache swapping overhead. + +Priority scheduling only takes effect with the TurboMind backend. The default scheduling policy is `fifo`; in that mode, requests are still scheduled by arrival order and the `priority` field does not change the scheduling order. + +## Enable Priority Scheduling + +Start the API Server with `--schedule-policy priority`: + +```bash +lmdeploy serve api_server internlm/internlm2_5-7b-chat \ + --backend turbomind \ + --schedule-policy priority +``` + +When creating a pipeline in Python, enable it through `TurbomindEngineConfig`: + +```python +from lmdeploy import GenerationConfig, TurbomindEngineConfig, pipeline + +backend_config = TurbomindEngineConfig(schedule_policy='priority') +pipe = pipeline('internlm/internlm2_5-7b-chat', backend_config=backend_config) + +response = pipe( + 'Introduce LMDeploy', + gen_config=GenerationConfig(priority=0, max_new_tokens=256), +) +``` + +`schedule_policy` supports the following values: + +| Value | Behavior | +| ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `fifo` | Default policy. Schedule requests by arrival order and preserve the original behavior. | +| `priority` | Schedule by request priority. Smaller values are admitted first; requests with the same priority keep FIFO order; already-started requests are kept before new requests. | + +## Set Request Priority + +Set request priority with `priority`. The valid range is `[0, 255]`. Smaller values have higher priority; `0` is the highest priority and `255` is the lowest priority. The default value is `0`. `priority` must be an integer; out-of-range values and strings, floats, booleans, or other non-integer types are rejected by validation. `null` is treated as not set (equivalent to `0`). + +### OpenAI-Compatible API + +Both `/v1/chat/completions` and `/v1/completions` support the LMDeploy extension field `priority`. + +```bash +curl http://localhost:23333/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "internlm/internlm2_5-7b-chat", + "messages": [{"role": "user", "content": "Summarize this text"}], + "priority": 0, + "max_tokens": 256 + }' +``` + +When using the OpenAI Python SDK, pass this extension field through `extra_body`: + +```python +from openai import OpenAI + +client = OpenAI(api_key='EMPTY', base_url='http://localhost:23333/v1') + +response = client.chat.completions.create( + model='internlm/internlm2_5-7b-chat', + messages=[{'role': 'user', 'content': 'Summarize this text'}], + max_tokens=256, + extra_body={'priority': 0}, +) +``` + +Use a larger value for normal-priority requests: + +```python +response = client.chat.completions.create( + model='internlm/internlm2_5-7b-chat', + messages=[{'role': 'user', 'content': 'Generate a long background report'}], + max_tokens=1024, + extra_body={'priority': 32}, +) +``` + +### Pipeline + +For pipeline calls, use `GenerationConfig.priority`: + +```python +from lmdeploy import GenerationConfig + +urgent = pipe( + 'Quickly review this code', + gen_config=GenerationConfig(priority=0, max_new_tokens=256), +) + +background = pipe( + 'Generate a long offline analysis report', + gen_config=GenerationConfig(priority=32, max_new_tokens=1024), +) +``` + +Priority scheduling is most visible when the service is handling concurrent requests. If there is only one synchronous request and no resource contention, changing its priority does not change the output content. + +## Scheduling Semantics + +After `schedule_policy='priority'` is enabled, TurboMind uses request priority in two stages: + +1. Before requests enter the engine, the waiting queue admits requests with smaller `priority` values first. +2. After requests enter the engine, already-started requests are kept first. Requests that have not started are then ordered by `priority` and arrival order. + +Therefore, `priority` is a non-preemptive priority policy: + +- A high-priority request can overtake lower-priority requests that are still waiting. +- A new high-priority request does not preempt a lower-priority request that has already started. +- Requests with the same priority are processed in arrival order. +- A continuous stream of high-priority requests may make lower-priority requests wait longer. The current policy does not include aging, deadlines, quotas, or weighted fair scheduling. + +This policy favors throughput and stability. It is a good fit for online services that need traffic differentiation while avoiding frequent preemption and KV cache swaps. + +## Usage Tips + +- Reserve smaller values, such as `0` or `1`, for the most urgent requests. +- Use middle values, such as `16` or `32`, for normal online requests. +- Use larger values, such as `128` or `255`, for background, offline, or delay-tolerant requests. +- If all requests use the default value `0`, scheduling is equivalent to FIFO among requests with the same priority. +- This feature only changes scheduling order. It does not change sampling parameters, output quality, or token generation rules. diff --git a/docs/en/index.rst b/docs/en/index.rst index 7063c84dde..647f7b712b 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -102,6 +102,7 @@ Documentation advance/long_context.md advance/chat_template.md advance/debug_turbomind.md + advance/turbomind_priority_scheduling.md advance/structed_output.md advance/pytorch_multinodes.md advance/pytorch_profiling.md diff --git a/docs/zh_cn/advance/turbomind_priority_scheduling.md b/docs/zh_cn/advance/turbomind_priority_scheduling.md new file mode 100644 index 0000000000..309d4fe5fe --- /dev/null +++ b/docs/zh_cn/advance/turbomind_priority_scheduling.md @@ -0,0 +1,125 @@ +# TurboMind 优先级调度 + +TurboMind 支持按请求优先级进行调度。该功能适合在同一个服务中同时处理多类请求，例如交互式请求、后台批处理任务、普通用户请求和高优先级用户请求。启用后，服务会优先让优先级更高的请求进入推理，同时尽量保留已经开始执行的请求，减少 KV cache 交换带来的额外开销。 + +优先级调度只在 TurboMind 后端生效。默认调度策略是 `fifo`，此时请求仍按到达顺序处理，`priority` 字段不会改变调度顺序。 + +## 启用优先级调度 + +启动 API Server 时设置 `--schedule-policy priority`： + +```bash +lmdeploy serve api_server internlm/internlm2_5-7b-chat \ + --backend turbomind \ + --schedule-policy priority +``` + +在 Python 中创建 pipeline 时，可以通过 `TurbomindEngineConfig` 启用： + +```python +from lmdeploy import GenerationConfig, TurbomindEngineConfig, pipeline + +backend_config = TurbomindEngineConfig(schedule_policy='priority') +pipe = pipeline('internlm/internlm2_5-7b-chat', backend_config=backend_config) + +response = pipe( + '介绍一下 LMDeploy', + gen_config=GenerationConfig(priority=0, max_new_tokens=256), +) +``` + +`schedule_policy` 支持以下取值： + +| 取值 | 行为 | +| ---------- | ----------------------------------------------------------------------------------------------------- | +| `fifo` | 默认策略。按请求到达顺序调度，保持原有行为。 | +| `priority` | 按请求优先级调度。数值更小的请求优先进入推理；同优先级请求保持 FIFO；已经开始执行的请求优先于新请求。 | + +## 设置请求优先级 + +请求优先级通过 `priority` 设置，取值范围为 `[0, 255]`。数值越小，优先级越高；`0` 是最高优先级，`255` 是最低优先级。未设置时默认值为 `0`。`priority` 必须是整数；超出范围或使用字符串、浮点数、布尔值等非整数类型会被校验拒绝。`null` 会被当作未设置（等价于 `0`）。 + +### OpenAI 兼容接口 + +`/v1/chat/completions` 和 `/v1/completions` 都支持 LMDeploy 扩展字段 `priority`。 + +```bash +curl http://localhost:23333/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "internlm/internlm2_5-7b-chat", + "messages": [{"role": "user", "content": "请总结这段文本"}], + "priority": 0, + "max_tokens": 256 + }' +``` + +使用 OpenAI Python SDK 时，可以通过 `extra_body` 传入该扩展字段： + +```python +from openai import OpenAI + +client = OpenAI(api_key='EMPTY', base_url='http://localhost:23333/v1') + +response = client.chat.completions.create( + model='internlm/internlm2_5-7b-chat', + messages=[{'role': 'user', 'content': '请总结这段文本'}], + max_tokens=256, + extra_body={'priority': 0}, +) +``` + +普通优先级请求可以使用更大的数值： + +```python +response = client.chat.completions.create( + model='internlm/internlm2_5-7b-chat', + messages=[{'role': 'user', 'content': '后台生成一份长报告'}], + max_tokens=1024, + extra_body={'priority': 32}, +) +``` + +### Pipeline + +Pipeline 调用中使用 `GenerationConfig.priority`： + +```python +from lmdeploy import GenerationConfig + +urgent = pipe( + '帮我快速检查这段代码', + gen_config=GenerationConfig(priority=0, max_new_tokens=256), +) + +background = pipe( + '生成一份较长的离线分析报告', + gen_config=GenerationConfig(priority=32, max_new_tokens=1024), +) +``` + +优先级调度主要在服务并发处理多个请求时体现效果。单个同步请求没有其他请求竞争资源时，设置不同优先级不会改变输出内容。 + +## 调度语义 + +启用 `schedule_policy='priority'` 后，TurboMind 会在两个阶段使用请求优先级： + +1. 请求进入 engine 前，等待队列优先取出 `priority` 数值更小的请求。 +2. 请求进入 engine 后，已经开始执行的请求优先保留；尚未开始执行的请求再按 `priority` 和到达顺序排序。 + +因此，`priority` 是非抢占式优先级策略： + +- 高优先级请求可以超过仍在等待的低优先级请求。 +- 高优先级新请求不会抢占已经开始执行的低优先级请求。 +- 同优先级请求按到达顺序处理。 +- 持续到来的高优先级请求可能让低优先级请求等待更久；当前策略不包含 aging、deadline、配额或加权公平调度。 + +这种策略更偏向吞吐和稳定性，适合希望区分请求等级，同时避免频繁抢占和 KV cache swap 的在线服务。 + +## 使用建议 + +- 为最紧急的请求保留较小的优先级，例如 `0` 或 `1`。 +- 为普通在线请求使用中间值，例如 `16` 或 `32`。 +- 为后台、离线或可延迟请求使用更大的值，例如 `128` 或 `255`。 +- 如果所有请求都使用默认值 `0`，调度效果等价于同优先级 FIFO。 +- 该功能只影响调度顺序，不改变采样参数、输出质量或 token 生成规则。 diff --git a/docs/zh_cn/index.rst b/docs/zh_cn/index.rst index c5b63b616d..55c8a38a92 100644 --- a/docs/zh_cn/index.rst +++ b/docs/zh_cn/index.rst @@ -103,6 +103,7 @@ LMDeploy 工具箱提供以下核心功能： advance/long_context.md advance/chat_template.md advance/debug_turbomind.md + advance/turbomind_priority_scheduling.md advance/structed_output.md advance/pytorch_multinodes.md advance/pytorch_profiling.md diff --git a/lmdeploy/cli/cli.py b/lmdeploy/cli/cli.py index 1dceba3c5a..fcc093dd7d 100644 --- a/lmdeploy/cli/cli.py +++ b/lmdeploy/cli/cli.py @@ -76,6 +76,7 @@ def add_parser_chat(): ArgumentHelper.communicator(tb_group) ArgumentHelper.cp(tb_group) ArgumentHelper.async_(tb_group) + ArgumentHelper.schedule_policy(tb_group) # speculative decoding ArgumentHelper.add_spec_group(parser) diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 9ca662b961..289472d12a 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -155,6 +155,7 @@ def add_parser_api_server(): ArgumentHelper.num_tokens_per_iter(tb_group) ArgumentHelper.max_prefill_iters(tb_group) ArgumentHelper.async_(tb_group) + ArgumentHelper.schedule_policy(tb_group) ArgumentHelper.communicator(tb_group) ArgumentHelper.dist_init_addr(tb_group) @@ -273,6 +274,7 @@ def api_server(args): num_tokens_per_iter=args.num_tokens_per_iter, max_prefill_iters=args.max_prefill_iters, async_=args.async_, + schedule_policy=args.schedule_policy, communicator=args.communicator, enable_metrics=not args.disable_metrics, hf_overrides=args.hf_overrides) diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 44805fe9d3..a871197f68 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -598,6 +598,16 @@ def async_(parser): help='Enable async execution (default: 1, enabled). ' 'Set to 0 to disable async mode, 1 to enable it.') + @staticmethod + def schedule_policy(parser): + return parser.add_argument('--schedule-policy', + type=str, + default='fifo', + choices=['fifo', 'priority'], + help='TurboMind scheduling policy. ' + '"fifo" preserves existing behavior; "priority" admits lower request priority ' + 'values first and avoids preempting started requests.') + @staticmethod def max_prefill_token_num(parser): return parser.add_argument('--max-prefill-token-num', diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index 9041486c38..4c268da2f8 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -108,6 +108,8 @@ class GenerationConfig: Must be non-negative; values below 0 are treated as 0. repetition_ngram_threshold: The number of times an n-gram must be repeated to trigger early stop. Must be non-negative; values below 0 are treated as 0. + priority: TurboMind scheduling priority. Smaller values have higher priority. + Must be in range [0, 255]. Defaults to 0. """ n: int = 1 @@ -146,6 +148,9 @@ class GenerationConfig: repetition_ngram_size: int = 0 repetition_ngram_threshold: int = 0 + # TurboMind scheduling priority. Smaller values have higher priority. + priority: int = 0 + def convert_stop_bad_words_to_ids(self, tokenizer: Tokenizer): """Convert stop_words/bad_sords to ids and append the ids to stop_token_ids/bad_token_ids.""" @@ -198,6 +203,8 @@ def __post_init__(self): if self.repetition_ngram_size <= 0 or self.repetition_ngram_threshold <= 0: self.repetition_ngram_size = 0 self.repetition_ngram_threshold = 0 + assert type(self.priority) is int, 'priority must be an integer' + assert 0 <= self.priority <= 255, 'priority must be in range [0, 255]' @pydantic_dataclass @@ -267,6 +274,10 @@ class TurbomindEngineConfig: hf_overrides: Huggingface overrides for the model. It can be used to override the default config of the model enable_metrics: enable metrics system + schedule_policy: TurboMind scheduling policy. + `fifo` preserves existing behavior. `priority` admits lower + GenerationConfig.priority values first and keeps already-started + requests before new requests inside TurboMind. """ dtype: str = 'auto' @@ -305,6 +316,7 @@ class TurbomindEngineConfig: communicator: str = 'nccl' hf_overrides: dict[str, Any] | None = None enable_metrics: bool = True + schedule_policy: Literal['fifo', 'priority'] = 'fifo' def __post_init__(self): """Check input validation.""" @@ -318,6 +330,7 @@ def __post_init__(self): 'invalid max_prefill_token_num' assert self.num_tokens_per_iter >= 0, 'invalid num_tokens_per_iter' assert self.async_ in (0, 1), 'async_ must be 0 (disabled) or 1 (enabled)' + assert self.schedule_policy in ('fifo', 'priority'), 'invalid schedule_policy' @dataclass diff --git a/lmdeploy/serve/openai/api_server.py b/lmdeploy/serve/openai/api_server.py index 17317fa23e..9522ace6d4 100644 --- a/lmdeploy/serve/openai/api_server.py +++ b/lmdeploy/serve/openai/api_server.py @@ -465,6 +465,7 @@ async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Reque logits_processors=logits_processors, min_new_tokens=request.min_new_tokens, min_p=request.min_p, + priority=request.priority or 0, random_seed=random_seed, spaces_between_special_tokens=request.spaces_between_special_tokens, migration_request=migration_request, @@ -760,6 +761,7 @@ async def completions_v1(request: CompletionRequest, raw_request: Request = None stop_words=request.stop, skip_special_tokens=request.skip_special_tokens, min_p=request.min_p, + priority=request.priority or 0, random_seed=random_seed, spaces_between_special_tokens=request.spaces_between_special_tokens, migration_request=migration_request, diff --git a/lmdeploy/serve/openai/protocol.py b/lmdeploy/serve/openai/protocol.py index 2f671737a0..ac53b2b628 100644 --- a/lmdeploy/serve/openai/protocol.py +++ b/lmdeploy/serve/openai/protocol.py @@ -5,7 +5,7 @@ from typing import Any, Literal import shortuuid -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, StrictInt class ErrorResponse(BaseModel): @@ -169,6 +169,7 @@ class ChatCompletionRequest(BaseModel): seed: int | None = None min_new_tokens: int | None = Field(default=None, examples=[None]) min_p: float = 0.0 + priority: StrictInt | None = Field(default=0, ge=0, le=255) enable_thinking: bool | None = None # will be deprecated in the future return_token_ids: bool | None = False include_stop_str_in_output: bool | None = False @@ -364,6 +365,7 @@ class CompletionRequest(BaseModel): top_k: int | None = 40 # for opencompass seed: int | None = None min_p: float = 0.0 + priority: StrictInt | None = Field(default=0, ge=0, le=255) class CompletionResponseChoice(BaseModel): diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py index b678e6c033..b9e4f1e569 100644 --- a/lmdeploy/turbomind/turbomind.py +++ b/lmdeploy/turbomind/turbomind.py @@ -251,6 +251,7 @@ def _from_hf(self, model_path: str, engine_config: TurbomindEngineConfig, ec.nnodes = engine_config.nnodes ec.node_rank = engine_config.node_rank ec.communicator = engine_config.communicator + ec.schedule_policy = engine_config.schedule_policy logger.info(f'turbomind engine config:\n\n' f'dtype={engine_config.dtype}, session_len={engine_config.session_len}, ' @@ -258,7 +259,8 @@ def _from_hf(self, model_path: str, engine_config: TurbomindEngineConfig, f'devices={engine_config.devices}, ' f'tp={engine_config.attn_tp_size}, ' f'dp={engine_config.attn_dp_size}, ' - f'cp={engine_config.attn_cp_size}') + f'cp={engine_config.attn_cp_size}, ' + f'schedule_policy={engine_config.schedule_policy}') model_comm = _tm.TurboMind.create(model_dir='', engine_config=ec) self._create_weight(model_comm) @@ -662,7 +664,7 @@ async def async_stream_infer(self, stream_output (bool): indicator for stream output kwargs (dict): kwargs for backward compatibility """ - logger.info(f'[async_stream_infer] session {session_id} start') + logger.info(f'[async_stream_infer] session {session_id} start, priority={gen_config.priority}') gen_cfg = self._get_generation_config(gen_config) inputs, input_len = self.prepare_inputs(input_ids=input_ids, @@ -791,6 +793,7 @@ def _get_generation_config(self, cfg: GenerationConfig): if not cfg.ignore_eos and cfg.stop_token_ids: c.stop_ids = _construct_stop_or_bad_words(cfg.stop_token_ids) c.repetition_penalty = cfg.repetition_penalty + c.priority = cfg.priority if cfg.min_new_tokens: c.min_new_tokens = cfg.min_new_tokens output_type = dict(all=1, generation=2) diff --git a/src/turbomind/engine/CMakeLists.txt b/src/turbomind/engine/CMakeLists.txt index 335aa99fac..bf0225d63e 100644 --- a/src/turbomind/engine/CMakeLists.txt +++ b/src/turbomind/engine/CMakeLists.txt @@ -13,3 +13,8 @@ add_library(engine STATIC target_link_libraries(engine PRIVATE xgrammar core) set_property(TARGET engine PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET engine PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) + +if(BUILD_TEST) + add_executable(test_request_queue test_request_queue.cc) + target_link_libraries(test_request_queue PRIVATE engine Catch2::Catch2WithMain) +endif() diff --git a/src/turbomind/engine/engine.cc b/src/turbomind/engine/engine.cc index 4bf9683c45..dab9b326dc 100644 --- a/src/turbomind/engine/engine.cc +++ b/src/turbomind/engine/engine.cc @@ -15,6 +15,7 @@ #include "src/turbomind/engine/engine.h" #include "src/turbomind/engine/model_executor.h" #include "src/turbomind/engine/request.h" +#include "src/turbomind/engine/schedule_policy.h" #include "src/turbomind/core/copy.h" #include "src/turbomind/core/logger.h" @@ -53,6 +54,26 @@ void serdes(Archive& ar, RequestData& r) ar& r.abort; } +namespace { + +uint64_t make_schedule_key(SchedulePolicy policy, bool scheduled, uint8_t priority, uint64_t arrival_order) +{ + switch (policy) { + case SchedulePolicy::kFifo: + return arrival_order; + case SchedulePolicy::kPriority: { + const uint64_t order = arrival_order & ((uint64_t{1} << 48) - 1); + const uint64_t scheduled_field = scheduled ? 0x00 : 0x0F; + return (scheduled_field << 56) | (uint64_t{priority} << 48) | order; + } + default: + TM_CHECK(false) << "unsupported schedule_policy: " << static_cast(policy); + return arrival_order; + } +} + +} // namespace + struct Engine::Impl { using Requests = vector>; @@ -109,6 +130,8 @@ struct Engine::Impl { const EngineParam param_; + const SchedulePolicy schedule_policy_; + Gateway& gateway_; comm::HostComm& tp_group_; @@ -186,6 +209,7 @@ Engine::Impl::Impl(EngineParam param, int queue_id, int phases): param_{param}, + schedule_policy_{parse_schedule_policy(param.schedule_policy)}, gateway_{gateway}, tp_group_{ctx.comm.h_tp_group}, dp_group_{ctx.comm.h_dp_group}, @@ -562,7 +586,8 @@ void Engine::Impl::Schedule() cache.push_back(c.get()); sequences.push_back(c->seq); status.push_back(c->seq->status); - priorities.push_back(c->req->unique_id); + priorities.push_back( + make_schedule_key(schedule_policy_, c->scheduled, c->req->gen_cfg.priority, c->req->unique_id)); context_length.push_back(c->seq_len + c->beta /* plus draft tokens */); alpha.push_back(c->alpha); TM_CHECK(c->seq->status == Sequence::kActive || c->alpha == 0) << c->seq->status << " " << c->alpha; @@ -609,6 +634,10 @@ void Engine::Impl::Schedule() // |<-- existing -->|<-- swap-in -->|<- swap-out ->| // |<----------- active ----------->|<------- inactive ----->| + for (auto i : swap_in) { + cache[i]->scheduled = true; + } + for (auto i : swap_in) { cache[i]->autoregres = {}; cache[i]->generating = {}; diff --git a/src/turbomind/engine/engine_config.h b/src/turbomind/engine/engine_config.h index 2c0381e9c4..76314dc9e5 100644 --- a/src/turbomind/engine/engine_config.h +++ b/src/turbomind/engine/engine_config.h @@ -27,6 +27,7 @@ struct EngineConfig { X(int, num_tokens_per_iter, 0) \ X(int, max_prefill_iters, 1) \ X(int, async_, 0) \ + X(std::string, schedule_policy, "fifo") \ X(int, outer_dp_size) \ X(int, attn_dp_size) \ X(int, attn_tp_size) \ diff --git a/src/turbomind/engine/gateway.cc b/src/turbomind/engine/gateway.cc index 82e3bd119a..b137522b71 100644 --- a/src/turbomind/engine/gateway.cc +++ b/src/turbomind/engine/gateway.cc @@ -8,11 +8,11 @@ namespace turbomind { -Gateway::Gateway(int size, std::function()> ctx_factory): - size_{size}, queues_(size_), dp_thr_{1}, ctx_factory_{ctx_factory}, next_{0} +Gateway::Gateway(int size, std::function()> ctx_factory, SchedulePolicy schedule_policy): + size_{size}, dp_thr_{1}, schedule_policy_{schedule_policy}, queues_(size_), ctx_factory_{ctx_factory}, next_{0} { for (int i = 0; i < size_; ++i) { - queues_[i] = std::make_unique(); + queues_[i] = RequestQueue::create(schedule_policy_); } signal_thread_ = std::thread(&Gateway::signal_thread_entry, this); diff --git a/src/turbomind/engine/gateway.h b/src/turbomind/engine/gateway.h index ddd80892d0..fd9860ff50 100644 --- a/src/turbomind/engine/gateway.h +++ b/src/turbomind/engine/gateway.h @@ -13,6 +13,7 @@ #include "src/turbomind/core/logger.h" #include "src/turbomind/engine/request.h" #include "src/turbomind/engine/request_queue.h" +#include "src/turbomind/engine/schedule_policy.h" #include "src/turbomind/engine/signal_buffer.h" namespace turbomind { @@ -62,7 +63,9 @@ class SequenceBinding { class Gateway { public: - Gateway(int size, std::function()> ctx_factory); + Gateway(int size, + std::function()> ctx_factory, + SchedulePolicy schedule_policy = SchedulePolicy::kFifo); void shutdown(); @@ -96,6 +99,8 @@ class Gateway { int dp_thr_; + SchedulePolicy schedule_policy_; + std::vector> queues_; std::function()> ctx_factory_; diff --git a/src/turbomind/engine/request.cc b/src/turbomind/engine/request.cc index f8bc0b28e0..f171c21304 100644 --- a/src/turbomind/engine/request.cc +++ b/src/turbomind/engine/request.cc @@ -36,6 +36,7 @@ std::ostream& operator<<(std::ostream& os, const GenerationConfig& c) os << ", temperature=" << c.temperature; os << ", repetition_penalty=" << c.repetition_penalty; os << ", random_seed=" << c.random_seed; + os << ", priority=" << static_cast(c.priority); os << ", output_logprobs=" << c.output_logprobs; os << ", output_hidden_states=" << c.output_last_hidden_state; os << ", output_logits=" << c.output_logits; diff --git a/src/turbomind/engine/request.h b/src/turbomind/engine/request.h index 1c1f226412..799baaec20 100644 --- a/src/turbomind/engine/request.h +++ b/src/turbomind/engine/request.h @@ -38,6 +38,8 @@ struct GenerationConfig { uint64_t random_seed = 0; + uint8_t priority = 0; + int output_logprobs = 0; enum OutType @@ -166,6 +168,8 @@ struct RequestCache { bool autoregres = false; // set at schedule, `seq_len` and `input_ids` taken from the engine bool generating = false; // set at schedule + bool scheduled = false; // set once when first selected as active + bool done = false; // set at cancel / update, is the request finished / canceled int alpha = 0; // pending growth of cache_len (draft_len + input_len) @@ -194,6 +198,7 @@ void serdes(Archive& ar, GenerationConfig& g) ar & g.temperature; ar & g.repetition_penalty; ar & g.random_seed; + ar & g.priority; ar & g.output_logprobs; ar & g.output_last_hidden_state; ar & g.output_logits; diff --git a/src/turbomind/engine/request_queue.cc b/src/turbomind/engine/request_queue.cc index b3e2c3aac3..9035cde276 100644 --- a/src/turbomind/engine/request_queue.cc +++ b/src/turbomind/engine/request_queue.cc @@ -7,4 +7,17 @@ namespace turbomind { +std::unique_ptr RequestQueue::create(SchedulePolicy schedule_policy) +{ + switch (schedule_policy) { + case SchedulePolicy::kFifo: + return std::make_unique(); + case SchedulePolicy::kPriority: + return std::make_unique(); + default: + throw std::invalid_argument("unsupported schedule_policy " + + std::to_string(static_cast(schedule_policy))); + } +} + } // namespace turbomind diff --git a/src/turbomind/engine/request_queue.h b/src/turbomind/engine/request_queue.h index 694f96bc79..199fc18750 100644 --- a/src/turbomind/engine/request_queue.h +++ b/src/turbomind/engine/request_queue.h @@ -2,18 +2,31 @@ #pragma once +#include +#include #include +#include #include +#include #include #include +#include +#include +#include +#include #include "src/turbomind/engine/request.h" +#include "src/turbomind/engine/schedule_policy.h" namespace turbomind { class RequestQueue { public: - explicit RequestQueue(): queue_{&pool_} {} + RequestQueue() = default; + + virtual ~RequestQueue() = default; + + static std::unique_ptr create(SchedulePolicy schedule_policy); void push(std::shared_ptr r) { @@ -22,7 +35,7 @@ class RequestQueue { if (closed_) { throw std::runtime_error("Queue is closed"); } - queue_.push_back(std::move(r)); + push_infer(std::move(r)); } cv_.notify_one(); } @@ -48,20 +61,14 @@ class RequestQueue { std::unique_lock lock{mutex_}; if (blocking) { - cv_.wait(lock, [this] { return !(queue_.empty() && kill_.empty()) || closed_; }); + cv_.wait(lock, [this] { return !(infer_queue_empty() && kill_.empty()) || closed_; }); } if (closed_) { abort = true; } - while (!queue_.empty() && infer_reqs.size() < max_infer) { - auto& r = queue_.front(); - if (r->cancel_flag.exchange(1, std::memory_order_acq_rel) == 0) { - infer_reqs.push_back(std::move(r)); - } - queue_.pop_front(); - } + pop_infer(infer_reqs, max_infer); kill_reqs.insert(kill_reqs.end(), kill_.begin(), kill_.end()); kill_.clear(); @@ -88,11 +95,19 @@ class RequestQueue { } } -private: +protected: + static bool try_claim_request(const std::shared_ptr& r) + { + return r->cancel_flag.exchange(1, std::memory_order_acq_rel) == 0; + } + + virtual bool infer_queue_empty() const = 0; + virtual void push_infer(std::shared_ptr r) = 0; + virtual void pop_infer(std::vector>& infer_reqs, unsigned max_infer) = 0; + std::atomic unique_id_{}; - std::pmr::unsynchronized_pool_resource pool_; - std::pmr::list> queue_; + std::pmr::unsynchronized_pool_resource pool_; std::vector> kill_; @@ -102,4 +117,77 @@ class RequestQueue { bool closed_{}; }; +class FifoRequestQueue: public RequestQueue { +public: + FifoRequestQueue(): queue_{&pool_} {} + +private: + bool infer_queue_empty() const override + { + return queue_.empty(); + } + + void push_infer(std::shared_ptr r) override + { + queue_.push_back(std::move(r)); + } + + void pop_infer(std::vector>& infer_reqs, unsigned max_infer) override + { + while (!queue_.empty() && infer_reqs.size() < max_infer) { + auto& r = queue_.front(); + if (try_claim_request(r)) { + infer_reqs.push_back(std::move(r)); + } + queue_.pop_front(); + } + } + + std::pmr::list> queue_; +}; + +class PriorityRequestQueue: public RequestQueue { +public: + PriorityRequestQueue(): queue_{&pool_} {} + +private: + struct PriorityEntry { + uint8_t priority; + uint64_t enqueue_order; + std::shared_ptr request; + }; + + static bool priority_entry_is_worse(const PriorityEntry& lhs, const PriorityEntry& rhs) + { + return std::tie(lhs.priority, lhs.enqueue_order) > std::tie(rhs.priority, rhs.enqueue_order); + } + + bool infer_queue_empty() const override + { + return queue_.empty(); + } + + void push_infer(std::shared_ptr r) override + { + const auto priority = r->gen_cfg.priority; + queue_.push_back(PriorityEntry{priority, enqueue_order_++, std::move(r)}); + std::push_heap(queue_.begin(), queue_.end(), priority_entry_is_worse); + } + + void pop_infer(std::vector>& infer_reqs, unsigned max_infer) override + { + while (!queue_.empty() && infer_reqs.size() < max_infer) { + std::pop_heap(queue_.begin(), queue_.end(), priority_entry_is_worse); + auto& entry = queue_.back(); + if (try_claim_request(entry.request)) { + infer_reqs.push_back(std::move(entry.request)); + } + queue_.pop_back(); + } + } + + uint64_t enqueue_order_{}; + std::pmr::vector queue_; +}; + } // namespace turbomind diff --git a/src/turbomind/engine/schedule_policy.h b/src/turbomind/engine/schedule_policy.h new file mode 100644 index 0000000000..8578ddfb60 --- /dev/null +++ b/src/turbomind/engine/schedule_policy.h @@ -0,0 +1,28 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#pragma once + +#include +#include +#include + +namespace turbomind { + +enum class SchedulePolicy +{ + kFifo, + kPriority, +}; + +inline SchedulePolicy parse_schedule_policy(std::string_view value) +{ + if (value == "fifo") { + return SchedulePolicy::kFifo; + } + if (value == "priority") { + return SchedulePolicy::kPriority; + } + throw std::invalid_argument{"invalid schedule_policy: " + std::string{value}}; +} + +} // namespace turbomind diff --git a/src/turbomind/engine/test_request_queue.cc b/src/turbomind/engine/test_request_queue.cc new file mode 100644 index 0000000000..ecabff1712 --- /dev/null +++ b/src/turbomind/engine/test_request_queue.cc @@ -0,0 +1,194 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#include +#include +#include +#include +#include +#include + +#include "catch2/catch_test_macros.hpp" + +#include "src/turbomind/engine/request_queue.h" + +namespace turbomind { +namespace { + +using Policy = SchedulePolicy; + +std::shared_ptr make_request(uint64_t id, uint8_t priority) +{ + auto req = std::make_shared(); + req->id = id; + req->session = SessionParam{id, 0, true, true, false}; + req->gen_cfg.priority = priority; + req->cancel_flag.store(0, std::memory_order_relaxed); + return req; +} + +void pop(RequestQueue& queue, + std::vector>& infer_reqs, + std::vector>& kill_reqs, + unsigned max_infer, + bool blocking = false) +{ + bool abort{}; + queue.pop(infer_reqs, kill_reqs, max_infer, blocking, abort); + REQUIRE_FALSE(abort); +} + +std::vector ids(const std::vector>& reqs) +{ + std::vector ret; + for (const auto& req : reqs) { + ret.push_back(req->id); + } + return ret; +} + +void expect_canceled_queued_request_is_skipped(RequestQueue& queue) +{ + auto canceled = make_request(1, 0); + auto valid = make_request(2, 5); + canceled->cancel_flag.store(-1, std::memory_order_relaxed); + queue.push(canceled); + queue.push(valid); + + std::vector> infer_reqs; + std::vector> kill_reqs; + pop(queue, infer_reqs, kill_reqs, 2); + + REQUIRE(ids(infer_reqs) == std::vector{2}); + REQUIRE(kill_reqs.empty()); +} + +} // namespace + +TEST_CASE("fifo policy ignores priority", "[request_queue]") +{ + FifoRequestQueue queue; + queue.push(make_request(1, 0)); + queue.push(make_request(2, 10)); + queue.push(make_request(3, 5)); + + std::vector> infer_reqs; + std::vector> kill_reqs; + pop(queue, infer_reqs, kill_reqs, 3); + + REQUIRE(ids(infer_reqs) == std::vector{1, 2, 3}); + REQUIRE(kill_reqs.empty()); +} + +TEST_CASE("priority policy returns smallest priority value first", "[request_queue]") +{ + PriorityRequestQueue queue; + queue.push(make_request(1, 0)); + queue.push(make_request(2, 10)); + queue.push(make_request(3, 5)); + + std::vector> infer_reqs; + std::vector> kill_reqs; + pop(queue, infer_reqs, kill_reqs, 3); + + REQUIRE(ids(infer_reqs) == std::vector{1, 3, 2}); + REQUIRE(kill_reqs.empty()); +} + +TEST_CASE("priority policy keeps fifo order for equal priority", "[request_queue]") +{ + PriorityRequestQueue queue; + queue.push(make_request(1, 3)); + queue.push(make_request(2, 3)); + queue.push(make_request(3, 3)); + + std::vector> infer_reqs; + std::vector> kill_reqs; + pop(queue, infer_reqs, kill_reqs, 3); + + REQUIRE(ids(infer_reqs) == std::vector{1, 2, 3}); +} + +TEST_CASE("priority policy max infer limits batch", "[request_queue]") +{ + PriorityRequestQueue queue; + queue.push(make_request(1, 1)); + queue.push(make_request(2, 9)); + queue.push(make_request(3, 8)); + + std::vector> infer_reqs; + std::vector> kill_reqs; + pop(queue, infer_reqs, kill_reqs, 2); + REQUIRE(ids(infer_reqs) == std::vector{1, 3}); + + infer_reqs.clear(); + kill_reqs.clear(); + pop(queue, infer_reqs, kill_reqs, 2); + REQUIRE(ids(infer_reqs) == std::vector{2}); +} + +TEST_CASE("canceled queued request is skipped in all policies", "[request_queue]") +{ + SECTION("fifo") + { + FifoRequestQueue queue; + expect_canceled_queued_request_is_skipped(queue); + } + + SECTION("priority") + { + PriorityRequestQueue queue; + expect_canceled_queued_request_is_skipped(queue); + } +} + +TEST_CASE("kill requests are drained independently", "[request_queue]") +{ + PriorityRequestQueue queue; + auto infer = make_request(1, 10); + auto kill = make_request(2, 0); + queue.push(infer); + queue.kill(kill); + + std::vector> infer_reqs; + std::vector> kill_reqs; + pop(queue, infer_reqs, kill_reqs, 1); + + REQUIRE(infer_reqs == std::vector>{infer}); + REQUIRE(kill_reqs == std::vector>{kill}); +} + +TEST_CASE("blocking pop wakes on push", "[request_queue]") +{ + PriorityRequestQueue queue; + + std::promise> promise; + auto future = promise.get_future(); + + std::thread worker{[&] { + std::vector> infer_reqs; + std::vector> kill_reqs; + pop(queue, infer_reqs, kill_reqs, 1, true); + promise.set_value(ids(infer_reqs)); + }}; + + std::this_thread::sleep_for(std::chrono::milliseconds{20}); + queue.push(make_request(42, 1)); + + REQUIRE(future.get() == std::vector{42}); + worker.join(); +} + +TEST_CASE("invalid schedule policy string is rejected", "[request_queue]") +{ + REQUIRE(parse_schedule_policy("fifo") == Policy::kFifo); + REQUIRE(parse_schedule_policy("priority") == Policy::kPriority); + REQUIRE_THROWS_AS(parse_schedule_policy("bad"), std::invalid_argument); +} + +TEST_CASE("request queue factory rejects unknown schedule policy", "[request_queue]") +{ + const auto invalid_policy = static_cast(0xFFFF); + REQUIRE_THROWS_AS(RequestQueue::create(invalid_policy), std::invalid_argument); +} + +} // namespace turbomind diff --git a/src/turbomind/python/bind.cpp b/src/turbomind/python/bind.cpp index 2b6b7faaf7..8d8cfd19a9 100644 --- a/src/turbomind/python/bind.cpp +++ b/src/turbomind/python/bind.cpp @@ -375,6 +375,7 @@ PYBIND11_MODULE(_turbomind, m) .def_readwrite("temperature", &ft::GenerationConfig::temperature) .def_readwrite("repetition_penalty", &ft::GenerationConfig::repetition_penalty) .def_readwrite("random_seed", &ft::GenerationConfig::random_seed) + .def_readwrite("priority", &ft::GenerationConfig::priority) .def_readwrite("output_logprobs", &ft::GenerationConfig::output_logprobs) .def_readwrite("output_last_hidden_state", &ft::GenerationConfig::output_last_hidden_state) .def_readwrite("output_logits", &ft::GenerationConfig::output_logits) diff --git a/src/turbomind/turbomind.cc b/src/turbomind/turbomind.cc index aac6512d9e..9350759373 100644 --- a/src/turbomind/turbomind.cc +++ b/src/turbomind/turbomind.cc @@ -15,6 +15,7 @@ #include "src/turbomind/engine/gateway.h" #include "src/turbomind/engine/model_executor.h" #include "src/turbomind/engine/model_request.h" +#include "src/turbomind/engine/schedule_policy.h" #include "src/turbomind/models/language_model.h" #include "src/turbomind/models/llama/context.h" @@ -47,6 +48,8 @@ struct TurboMind::Impl { unique_ptr group_id_; + SchedulePolicy schedule_policy_; + shared_ptr gateway_; FFICtxFactory ffi_ctx_factory_; @@ -146,7 +149,10 @@ TurboMind::Impl::~Impl() } TurboMind::Impl::Impl(string model_dir, EngineConfig config, FFICtxFactory ffi_ctx_factory): - data_type_{}, engine_param_{}, ffi_ctx_factory_{ffi_ctx_factory} + data_type_{}, + engine_param_{}, + schedule_policy_{parse_schedule_policy(config.schedule_policy)}, + ffi_ctx_factory_{ffi_ctx_factory} { data_type_ = config.data_type; TM_CHECK(data_type_ == kBfloat16 || data_type_ == kHalf); @@ -255,7 +261,7 @@ void TurboMind::Impl::CreateContext(int index) for (size_t i = 0; i < queue_id_.size(); ++i) { queue_id_[i] = queue_id_[i] ? n_queues_++ : -1; } - gateway_ = std::make_shared(n_queues_, ffi_ctx_factory_); + gateway_ = std::make_shared(n_queues_, ffi_ctx_factory_, schedule_policy_); } c.h_global->Sync(); diff --git a/tests/test_lmdeploy/test_messages.py b/tests/test_lmdeploy/test_messages.py index cb57079f07..4c48f9ecae 100644 --- a/tests/test_lmdeploy/test_messages.py +++ b/tests/test_lmdeploy/test_messages.py @@ -2,7 +2,7 @@ import pytest from pydantic import ValidationError -from lmdeploy import GenerationConfig, Tokenizer +from lmdeploy import GenerationConfig, Tokenizer, TurbomindEngineConfig from lmdeploy.serve.openai.protocol import ChatCompletionRequest from lmdeploy.utils import get_hf_gen_cfg @@ -13,6 +13,29 @@ def test_generation_config_repetition_ngram_clamped(): assert c.repetition_ngram_threshold == 0 +def test_generation_config_priority_validation(): + assert GenerationConfig(priority=0).priority == 0 + assert GenerationConfig(priority=7).priority == 7 + assert GenerationConfig(priority=255).priority == 255 + + for priority in (None, True, 1.5, '1'): + with pytest.raises(AssertionError): + GenerationConfig(priority=priority) + + for priority in (-1, 256): + with pytest.raises(AssertionError): + GenerationConfig(priority=priority) + + +def test_turbomind_engine_config_schedule_policy_validation(): + assert TurbomindEngineConfig().schedule_policy == 'fifo' + assert TurbomindEngineConfig(schedule_policy='fifo').schedule_policy == 'fifo' + assert TurbomindEngineConfig(schedule_policy='priority').schedule_policy == 'priority' + + with pytest.raises((AssertionError, ValidationError)): + TurbomindEngineConfig(schedule_policy='bad') + + def test_chat_completion_request_repetition_ngram_ge_zero(): with pytest.raises(ValidationError): ChatCompletionRequest( diff --git a/tests/test_lmdeploy/test_openai_protocol_priority.py b/tests/test_lmdeploy/test_openai_protocol_priority.py new file mode 100644 index 0000000000..f9b0ac95c6 --- /dev/null +++ b/tests/test_lmdeploy/test_openai_protocol_priority.py @@ -0,0 +1,26 @@ +import pytest +from pydantic import ValidationError + +from lmdeploy.serve.openai.protocol import ChatCompletionRequest, CompletionRequest + + +def test_chat_completion_request_priority_extension(): + assert ChatCompletionRequest(model='m', messages=[{'role': 'user', 'content': 'hi'}]).priority == 0 + assert ChatCompletionRequest(model='m', messages=[{'role': 'user', 'content': 'hi'}], + priority=255).priority == 255 + assert ChatCompletionRequest(model='m', messages=[{'role': 'user', 'content': 'hi'}], + priority=None).priority is None + + for priority in (-1, 256, True, 1.5, '1'): + with pytest.raises(ValidationError): + ChatCompletionRequest(model='m', messages=[{'role': 'user', 'content': 'hi'}], priority=priority) + + +def test_completion_request_priority_extension(): + assert CompletionRequest(model='m', prompt='hi').priority == 0 + assert CompletionRequest(model='m', prompt='hi', priority=5).priority == 5 + assert CompletionRequest(model='m', prompt='hi', priority=None).priority is None + + for priority in (-1, 256, True, 1.5, '1'): + with pytest.raises(ValidationError): + CompletionRequest(model='m', prompt='hi', priority=priority)