mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2026-04-23 00:17:25 +08:00
[Others] reuse unit test (#7127)
This commit is contained in:
luukunn
@@ -0,0 +1,82 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from fastdeploy.cache_manager.multimodal_cache_manager import EncoderCacheManager
|
||||
from fastdeploy.engine.request import ImagePosition
|
||||
|
||||
|
||||
def test_mm_encoder_cache():
|
||||
max_encoder_cache = 4096
|
||||
encoder_cache = EncoderCacheManager(max_encoder_cache=max_encoder_cache)
|
||||
|
||||
mm_hashes = ["mm_hash1", "mm_hash2"]
|
||||
mm_positions = [ImagePosition(offset=120, length=400), ImagePosition(offset=620, length=800)]
|
||||
|
||||
cache_length = mm_positions[0].length + mm_positions[1].length
|
||||
evict_hashes = encoder_cache.apply_cache(mm_hashes=mm_hashes, mm_items=mm_positions)
|
||||
assert evict_hashes == [], "The evicted hashes should be empty."
|
||||
assert list(encoder_cache.cache.keys()) == [
|
||||
"mm_hash1",
|
||||
"mm_hash2",
|
||||
], "The cache should contain mm_hash1 and mm_hash2."
|
||||
assert (
|
||||
encoder_cache.current_cache_size == cache_length
|
||||
), "The cache size should be the sum of the lengths of mm_hash1 and mm_hash2."
|
||||
assert (
|
||||
encoder_cache.current_cache_size <= max_encoder_cache
|
||||
), "The cache size should be less than or equal to the max_encoder_cache."
|
||||
|
||||
mm_hashes = ["mm_hash3", "mm_hash4"]
|
||||
mm_positions = [ImagePosition(offset=20, length=1204), ImagePosition(offset=1800, length=2048)]
|
||||
cache_length += mm_positions[0].length + mm_positions[1].length - 400
|
||||
evict_hashes = encoder_cache.apply_cache(mm_hashes=mm_hashes, mm_items=mm_positions)
|
||||
assert evict_hashes == ["mm_hash1"], "The evicted hashes should be mm_hash1."
|
||||
assert list(encoder_cache.cache.keys()) == [
|
||||
"mm_hash2",
|
||||
"mm_hash3",
|
||||
"mm_hash4",
|
||||
], "The cache should contain mm_hash2, mm_hash3, and mm_hash4."
|
||||
assert (
|
||||
encoder_cache.current_cache_size == cache_length
|
||||
), "The cache size should be the sum of the lengths of mm_hash2, mm_hash3, and mm_hash4."
|
||||
assert (
|
||||
encoder_cache.current_cache_size <= max_encoder_cache
|
||||
), "The cache size should be less than or equal to the max_encoder_cache."
|
||||
|
||||
evict_hashes = encoder_cache.apply_cache(mm_hashes=["mm_hash2"], mm_items=[ImagePosition(offset=620, length=800)])
|
||||
assert evict_hashes == [], "The evicted hashes should be empty."
|
||||
assert (
|
||||
encoder_cache.current_cache_size == cache_length
|
||||
), "The cache size should be the sum of the lengths of mm_hash2, mm_hash3, and mm_hash4."
|
||||
assert (
|
||||
encoder_cache.current_cache_size <= max_encoder_cache
|
||||
), "The cache size should be less than or equal to the max_encoder_cache."
|
||||
|
||||
cache_length -= 1204
|
||||
evict_hashes = encoder_cache.evict_cache(needed=800)
|
||||
assert evict_hashes == ["mm_hash3"], "The evicted hashes should be mm_hash3."
|
||||
assert list(encoder_cache.cache.keys()) == [
|
||||
"mm_hash4",
|
||||
"mm_hash2",
|
||||
], "The cache should contain mm_hash2 and mm_hash4."
|
||||
assert (
|
||||
encoder_cache.current_cache_size == cache_length
|
||||
), "The cache size should be the sum of the lengths of mm_hash2 and mm_hash4."
|
||||
assert (
|
||||
encoder_cache.current_cache_size <= max_encoder_cache
|
||||
), "The cache size should be less than or equal to the max_encoder_cache."
|
||||
|
||||
encoder_cache.clear_cache()
|
||||
assert encoder_cache.current_cache_size == 0, "The cache size should be 0."
|
||||
assert list(encoder_cache.cache.keys()) == [], "The cache should be empty."
|
||||
@@ -0,0 +1,270 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import asdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastdeploy.cache_manager.prefix_cache_manager import PrefixCacheManager
|
||||
from fastdeploy.config import CacheConfig, FDConfig, ParallelConfig
|
||||
from fastdeploy.engine.args_utils import EngineArgs
|
||||
from fastdeploy.engine.request import ImagePosition, Request
|
||||
from fastdeploy.scheduler import SchedulerConfig
|
||||
|
||||
|
||||
def make_prefix_cache_manager(max_num_seqs, enable_mm=False, num_gpu_blocks_override=100, max_num_batched_tokens=3200):
|
||||
engine_args = EngineArgs(
|
||||
max_num_seqs=max_num_seqs,
|
||||
num_gpu_blocks_override=num_gpu_blocks_override,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
)
|
||||
args = asdict(engine_args)
|
||||
cache_cfg = CacheConfig(args)
|
||||
model_cfg = SimpleNamespace(enable_mm=enable_mm, max_model_len=4196)
|
||||
speculative_cfg = SimpleNamespace(method=None)
|
||||
model_cfg.print = print
|
||||
model_cfg.architectures = ["test_model"]
|
||||
model_cfg.mm_max_tokens_per_item = None
|
||||
model_cfg.version = None # Required for register_info
|
||||
cache_cfg.bytes_per_token_per_layer = 1
|
||||
|
||||
parallel_cfg = ParallelConfig(args)
|
||||
scheduler_cfg = SchedulerConfig(args)
|
||||
graph_opt_cfg = engine_args.create_graph_optimization_config()
|
||||
fd_config = FDConfig(
|
||||
model_config=model_cfg,
|
||||
cache_config=cache_cfg,
|
||||
parallel_config=parallel_cfg,
|
||||
graph_opt_config=graph_opt_cfg,
|
||||
speculative_config=speculative_cfg,
|
||||
scheduler_config=scheduler_cfg,
|
||||
)
|
||||
return PrefixCacheManager(config=fd_config, tensor_parallel_size=8, splitwise_role="mixed")
|
||||
|
||||
|
||||
def test_block_num_limit():
|
||||
import pytest
|
||||
|
||||
with pytest.raises(AssertionError):
|
||||
make_prefix_cache_manager(max_num_seqs=3, enable_mm=False, num_gpu_blocks_override=20)
|
||||
|
||||
|
||||
def test_normal_case():
|
||||
block_size = 64
|
||||
cache_manager = make_prefix_cache_manager(max_num_seqs=3, enable_mm=False, num_gpu_blocks_override=128)
|
||||
req1 = Request.from_dict({"request_id": "req1", "prompt_token_ids": [1] * 3200, "prompt_token_ids_len": 3200})
|
||||
req2 = Request.from_dict(
|
||||
{"request_id": "req2", "prompt_token_ids": [1] * 1600 + [2] * 1600, "prompt_token_ids_len": 3200}
|
||||
)
|
||||
req3 = Request.from_dict(
|
||||
{"request_id": "req3", "prompt_token_ids": [1] * 1600 + [3] * 1600, "prompt_token_ids_len": 3200}
|
||||
)
|
||||
(common_block_ids, matched_token_num, hit_info) = cache_manager.request_match_blocks(req1, block_size)
|
||||
assert len(common_block_ids) == 0
|
||||
assert matched_token_num == 0
|
||||
assert len(cache_manager.gpu_free_block_list) == 128
|
||||
req1.block_tables.extend(common_block_ids)
|
||||
# allocate for req1 inputs
|
||||
num_new_block = 50
|
||||
req1.block_tables.extend(cache_manager.allocate_gpu_blocks(num_new_block))
|
||||
req1.num_computed_tokens += 50 * block_size
|
||||
cache_manager.update_cache_blocks(req1, block_size, req1.num_computed_tokens)
|
||||
assert len(cache_manager.gpu_free_block_list) == 78
|
||||
# allocate for req2 inputs
|
||||
(common_block_ids, matched_token_num, hit_info) = cache_manager.request_match_blocks(req2, block_size)
|
||||
assert len(common_block_ids) == 25
|
||||
assert matched_token_num == 25 * block_size
|
||||
req2.num_cached_tokens = matched_token_num
|
||||
req2.num_computed_tokens = 25 * block_size
|
||||
num_new_block = 25
|
||||
req2.block_tables.extend(common_block_ids)
|
||||
req2.block_tables.extend(cache_manager.allocate_gpu_blocks(num_new_block))
|
||||
cache_manager.update_cache_blocks(req2, block_size, req2.num_computed_tokens)
|
||||
# allocate for req3 input
|
||||
(common_block_ids, matched_token_num, hit_info) = cache_manager.request_match_blocks(req3, block_size)
|
||||
assert len(common_block_ids) == 25
|
||||
assert matched_token_num == 25 * block_size
|
||||
req3.num_cached_tokens = matched_token_num
|
||||
req3.num_computed_tokens = 25 * block_size
|
||||
assert len(cache_manager.gpu_free_block_list) == 53
|
||||
req3.block_tables.extend(common_block_ids)
|
||||
num_new_block = 25
|
||||
assert cache_manager.can_allocate_gpu_blocks(num_new_block)
|
||||
req3.block_tables.extend(cache_manager.allocate_gpu_blocks(num_new_block))
|
||||
cache_manager.update_cache_blocks(req3, block_size, req3.num_computed_tokens)
|
||||
assert len(cache_manager.gpu_free_block_list) == 28
|
||||
|
||||
|
||||
def test_mm_extra_keys():
|
||||
block_size = 64
|
||||
cache_manager = make_prefix_cache_manager(max_num_seqs=3, enable_mm=True)
|
||||
|
||||
prompt_token_ids = [1] * 100 + [2] * 100
|
||||
req1 = {
|
||||
"request_id": "req1",
|
||||
"prompt_token_ids": prompt_token_ids,
|
||||
"prompt_token_ids_len": len(prompt_token_ids),
|
||||
}
|
||||
for idx in range(0, len(prompt_token_ids), block_size):
|
||||
token_ids_lens = min(block_size, len(prompt_token_ids[idx:]))
|
||||
mm_idx, extra_keys = cache_manager.get_block_hash_extra_keys(
|
||||
request=Request.from_dict(req1),
|
||||
start_idx=idx,
|
||||
end_idx=idx + token_ids_lens,
|
||||
mm_idx=0,
|
||||
)
|
||||
assert extra_keys == [], f"extra_keys {extra_keys} != [], start_idx {idx}, end_idx {idx + token_ids_lens}"
|
||||
assert mm_idx == 0, f"mm_idx {mm_idx} != 0, start_idx {idx}, end_idx {idx + token_ids_lens}"
|
||||
|
||||
# block 1
|
||||
prompt_token_ids = [1] * 30 + [-1] * 34
|
||||
mm_positions = [ImagePosition(offset=30, length=80)]
|
||||
mm_hashes = ["image1"]
|
||||
extra_keys_list = [(0, ["image1"])]
|
||||
|
||||
# block 2
|
||||
prompt_token_ids += [-1] * 46 + [2] * 18
|
||||
extra_keys_list.append((1, ["image1"]))
|
||||
|
||||
# block 3
|
||||
prompt_token_ids += [-1] * 100
|
||||
mm_positions.append(ImagePosition(offset=128, length=100))
|
||||
mm_hashes.append("image2")
|
||||
extra_keys_list.append((1, ["image2"]))
|
||||
|
||||
# block 4、5
|
||||
prompt_token_ids += [3] * 40
|
||||
extra_keys_list.append((1, ["image2"]))
|
||||
extra_keys_list.append((1, []))
|
||||
|
||||
req2 = {
|
||||
"request_id": "req2",
|
||||
"prompt_token_ids": prompt_token_ids,
|
||||
"prompt_token_ids_len": len(prompt_token_ids),
|
||||
"multimodal_inputs": {
|
||||
"mm_positions": mm_positions,
|
||||
"mm_hashes": mm_hashes,
|
||||
},
|
||||
}
|
||||
|
||||
mm_idx, key_idx = 0, 0
|
||||
for idx in range(0, len(prompt_token_ids), block_size):
|
||||
token_ids_lens = min(block_size, len(prompt_token_ids[idx:]))
|
||||
mm_idx, extra_keys = cache_manager.get_block_hash_extra_keys(
|
||||
request=Request.from_dict(req2),
|
||||
start_idx=idx,
|
||||
end_idx=idx + token_ids_lens,
|
||||
mm_idx=mm_idx,
|
||||
)
|
||||
|
||||
target_idx, target_keys = extra_keys_list[key_idx]
|
||||
assert (
|
||||
mm_idx == target_idx
|
||||
), f"mm_idx {mm_idx} != target_idx {target_idx}, start_idx {idx}, end_idx {idx + token_ids_lens}"
|
||||
assert (
|
||||
extra_keys == target_keys
|
||||
), f"extra_keys {extra_keys} != target_keys {target_keys}, start_idx {idx}, end_idx {idx + token_ids_lens}"
|
||||
key_idx += 1
|
||||
|
||||
|
||||
def test_mm_prefix_cache():
|
||||
block_size = 64
|
||||
cache_manager = make_prefix_cache_manager(max_num_seqs=3, enable_mm=True, num_gpu_blocks_override=100)
|
||||
multimodal_inputs = {
|
||||
"mm_positions": [ImagePosition(offset=120, length=1200)],
|
||||
"mm_hashes": ["image1"],
|
||||
}
|
||||
req1_dict = {
|
||||
"request_id": "req1",
|
||||
"prompt_token_ids": [1] * 120 + [-1] * 1200 + [2] * 120,
|
||||
"prompt_token_ids_len": 1440,
|
||||
"multimodal_inputs": multimodal_inputs,
|
||||
}
|
||||
req1 = Request.from_dict(req1_dict)
|
||||
|
||||
multimodal_inputs = dict(multimodal_inputs)
|
||||
multimodal_inputs["mm_positions"].append(ImagePosition(offset=1836, length=587))
|
||||
multimodal_inputs["mm_hashes"].append("image2")
|
||||
req2_dict = {
|
||||
"request_id": "req2",
|
||||
"prompt_token_ids": [1] * 120 + [-1] * 1200 + [2] * 120 + [3] * 396 + [-1] * 587,
|
||||
"prompt_token_ids_len": 2423,
|
||||
"multimodal_inputs": multimodal_inputs,
|
||||
}
|
||||
req2 = Request.from_dict(req2_dict)
|
||||
|
||||
multimodal_inputs = dict(multimodal_inputs)
|
||||
multimodal_inputs["mm_hashes"] = ["image3", "image4"]
|
||||
req3_dict = {
|
||||
"request_id": "req3",
|
||||
"prompt_token_ids": [1] * 120 + [-1] * 1200 + [2] * 120 + [3] * 396 + [-1] * 587,
|
||||
"prompt_token_ids_len": 2423,
|
||||
"multimodal_inputs": multimodal_inputs,
|
||||
}
|
||||
req3 = Request.from_dict(req3_dict)
|
||||
|
||||
multimodal_inputs = dict(multimodal_inputs)
|
||||
multimodal_inputs["mm_positions"] = [ImagePosition(offset=120, length=1200)]
|
||||
multimodal_inputs["mm_hashes"] = ["image3"]
|
||||
req4_dict = {
|
||||
"request_id": "req4",
|
||||
"prompt_token_ids": [1] * 120 + [-1] * 1200 + [2] * 120 + [3] * 352,
|
||||
"prompt_token_ids_len": 1792,
|
||||
"multimodal_inputs": multimodal_inputs,
|
||||
}
|
||||
req4 = Request.from_dict(req4_dict)
|
||||
|
||||
(common_block_ids, matched_token_num, hit_info) = cache_manager.request_match_blocks(req1, block_size)
|
||||
assert len(common_block_ids) == 0
|
||||
assert matched_token_num == 0
|
||||
assert len(cache_manager.gpu_free_block_list) == 100
|
||||
req1.block_tables.extend(common_block_ids)
|
||||
|
||||
# allocate for req1 inputs
|
||||
num_new_block = 22
|
||||
req1.block_tables.extend(cache_manager.allocate_gpu_blocks(num_new_block))
|
||||
req1.num_computed_tokens += 22 * block_size
|
||||
cache_manager.update_cache_blocks(req1, block_size, req1.num_computed_tokens)
|
||||
assert len(cache_manager.gpu_free_block_list) == 78
|
||||
|
||||
# allocate for req2 inputs
|
||||
(common_block_ids, matched_token_num, hit_info) = cache_manager.request_match_blocks(req2, block_size)
|
||||
assert len(common_block_ids) == 22
|
||||
assert matched_token_num == 22 * block_size
|
||||
req2.num_cached_tokens = matched_token_num
|
||||
req2.num_computed_tokens = matched_token_num
|
||||
num_new_block = 15
|
||||
req2.block_tables.extend(common_block_ids)
|
||||
req2.block_tables.extend(cache_manager.allocate_gpu_blocks(num_new_block))
|
||||
req2.num_computed_tokens += 15 * block_size
|
||||
cache_manager.update_cache_blocks(req2, block_size, req2.num_computed_tokens)
|
||||
|
||||
# allocate for req3 input
|
||||
(common_block_ids, matched_token_num, hit_info) = cache_manager.request_match_blocks(req3, block_size)
|
||||
assert len(common_block_ids) == 1
|
||||
assert matched_token_num == 1 * block_size
|
||||
req3.num_cached_tokens = matched_token_num
|
||||
req3.num_computed_tokens = matched_token_num
|
||||
assert len(cache_manager.gpu_free_block_list) == 63
|
||||
req3.block_tables.extend(common_block_ids)
|
||||
num_new_block = 36
|
||||
assert cache_manager.can_allocate_gpu_blocks(num_new_block)
|
||||
req3.block_tables.extend(cache_manager.allocate_gpu_blocks(num_new_block))
|
||||
req3.num_computed_tokens += 36 * block_size
|
||||
cache_manager.update_cache_blocks(req3, block_size, req3.num_computed_tokens)
|
||||
assert len(cache_manager.gpu_free_block_list) == 27
|
||||
|
||||
# allocate for req4 input
|
||||
(common_block_ids, matched_token_num, hit_info) = cache_manager.request_match_blocks(req4, block_size)
|
||||
assert len(common_block_ids) == 28
|
||||
assert matched_token_num == 28 * block_size
|
||||
@@ -0,0 +1,728 @@
|
||||
"""
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
|
||||
import concurrent.futures
|
||||
import pickle
|
||||
import unittest
|
||||
from dataclasses import asdict
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import paddle
|
||||
|
||||
if not hasattr(paddle, "compat"):
|
||||
paddle.compat = SimpleNamespace(enable_torch_proxy=lambda scope: None)
|
||||
|
||||
from fastdeploy.config import CacheConfig, FDConfig, ParallelConfig, SchedulerConfig
|
||||
from fastdeploy.engine.args_utils import EngineArgs
|
||||
from fastdeploy.engine.request import (
|
||||
CompletionOutput,
|
||||
ImagePosition,
|
||||
Request,
|
||||
RequestMetrics,
|
||||
RequestOutput,
|
||||
RequestStatus,
|
||||
)
|
||||
from fastdeploy.engine.sched.resource_manager_v1 import (
|
||||
ResourceManagerV1,
|
||||
SignalConsumer,
|
||||
)
|
||||
from fastdeploy.input.utils import IDS_TYPE_FLAG
|
||||
|
||||
|
||||
def _build_manager(
|
||||
splitwise_role="mixed",
|
||||
enable_mm=True,
|
||||
enable_prefix_caching=True,
|
||||
disable_chunked_mm_input=False,
|
||||
speculative_method=None,
|
||||
block_size=4,
|
||||
max_num_batched_tokens=128,
|
||||
max_model_len=64,
|
||||
architectures=None,
|
||||
max_encoder_cache=0,
|
||||
max_processor_cache=0,
|
||||
num_gpu_blocks_override=128,
|
||||
):
|
||||
max_num_seqs = 2
|
||||
engine_args = EngineArgs(
|
||||
max_num_seqs=max_num_seqs,
|
||||
num_gpu_blocks_override=num_gpu_blocks_override,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
)
|
||||
args = asdict(engine_args)
|
||||
|
||||
cache_cfg = CacheConfig(args)
|
||||
cache_cfg.block_size = block_size
|
||||
cache_cfg.max_block_num_per_seq = 8
|
||||
cache_cfg.enc_dec_block_num = 1
|
||||
cache_cfg.enable_prefix_caching = enable_prefix_caching
|
||||
cache_cfg.enable_output_caching = True
|
||||
cache_cfg.disable_chunked_mm_input = disable_chunked_mm_input
|
||||
cache_cfg.max_encoder_cache = max_encoder_cache
|
||||
cache_cfg.max_processor_cache = max_processor_cache
|
||||
model_cfg = SimpleNamespace(enable_mm=enable_mm)
|
||||
speculative_cfg = SimpleNamespace(method=speculative_method, num_speculative_tokens=1)
|
||||
model_cfg.print = print
|
||||
model_cfg.max_model_len = max_model_len
|
||||
model_cfg.architectures = architectures or ["test_model"]
|
||||
model_cfg.mm_max_tokens_per_item = None
|
||||
model_cfg.version = None # Required for register_info
|
||||
cache_cfg.bytes_per_token_per_layer = 1
|
||||
cache_cfg.kv_cache_ratio = 1.0
|
||||
parallel_cfg = ParallelConfig(args)
|
||||
scheduler_cfg = SchedulerConfig(args)
|
||||
scheduler_cfg.splitwise_role = splitwise_role
|
||||
graph_opt_cfg = engine_args.create_graph_optimization_config()
|
||||
|
||||
fd_config = FDConfig(
|
||||
model_config=model_cfg,
|
||||
cache_config=cache_cfg,
|
||||
parallel_config=parallel_cfg,
|
||||
graph_opt_config=graph_opt_cfg,
|
||||
speculative_config=speculative_cfg,
|
||||
scheduler_config=scheduler_cfg,
|
||||
)
|
||||
return ResourceManagerV1(
|
||||
max_num_seqs=max_num_seqs,
|
||||
config=fd_config,
|
||||
tensor_parallel_size=2,
|
||||
splitwise_role=splitwise_role,
|
||||
)
|
||||
|
||||
|
||||
def _make_request(request_id="req-1", prompt_token_ids=None, multimodal_inputs=None):
|
||||
req_dict = {
|
||||
"request_id": request_id,
|
||||
"multimodal_inputs": multimodal_inputs or {},
|
||||
}
|
||||
request = Request.from_dict(req_dict)
|
||||
request.prompt_token_ids = prompt_token_ids or [1, 2, 3, 4]
|
||||
request.prompt_token_ids_len = len(request.prompt_token_ids)
|
||||
request.need_prefill_tokens = request.prompt_token_ids_len
|
||||
request.output_token_ids = []
|
||||
request.disaggregate_info = {}
|
||||
request.metrics = RequestMetrics()
|
||||
return request
|
||||
|
||||
|
||||
def _register_manager_cleanup(testcase, manager):
|
||||
testcase.addCleanup(manager.need_block_num_signal.clear)
|
||||
testcase.addCleanup(manager.finish_execution_pool.shutdown, wait=True)
|
||||
testcase.addCleanup(manager.async_preprocess_pool.shutdown, wait=True)
|
||||
|
||||
|
||||
class TestResourceManagerV1(unittest.TestCase):
|
||||
def setUp(self):
|
||||
max_num_seqs = 2
|
||||
engine_args = EngineArgs(
|
||||
max_num_seqs=max_num_seqs,
|
||||
num_gpu_blocks_override=102,
|
||||
max_num_batched_tokens=3200,
|
||||
)
|
||||
args = asdict(engine_args)
|
||||
|
||||
cache_cfg = CacheConfig(args)
|
||||
model_cfg = SimpleNamespace(enable_mm=True) # Enable multimodal for feature testing
|
||||
speculative_cfg = SimpleNamespace(method=None)
|
||||
model_cfg.print = print
|
||||
model_cfg.max_model_len = 3200
|
||||
model_cfg.architectures = ["test_model"]
|
||||
model_cfg.mm_max_tokens_per_item = None
|
||||
model_cfg.version = None # Required for register_info
|
||||
cache_cfg.bytes_per_token_per_layer = 1
|
||||
cache_cfg.kv_cache_ratio = 1.0
|
||||
parallel_cfg = ParallelConfig(args)
|
||||
scheduler_cfg = SchedulerConfig(args)
|
||||
graph_opt_cfg = engine_args.create_graph_optimization_config()
|
||||
|
||||
fd_config = FDConfig(
|
||||
model_config=model_cfg,
|
||||
cache_config=cache_cfg,
|
||||
parallel_config=parallel_cfg,
|
||||
graph_opt_config=graph_opt_cfg,
|
||||
speculative_config=speculative_cfg,
|
||||
scheduler_config=scheduler_cfg,
|
||||
)
|
||||
self.manager = ResourceManagerV1(
|
||||
max_num_seqs=max_num_seqs, config=fd_config, tensor_parallel_size=8, splitwise_role="mixed"
|
||||
)
|
||||
req_dict = {
|
||||
"request_id": "test_request",
|
||||
"multimodal_inputs": {},
|
||||
}
|
||||
self.request = Request.from_dict(req_dict)
|
||||
self.request.async_process_futures = []
|
||||
self.request.multimodal_inputs = {}
|
||||
|
||||
def test_waiting_async_process_no_futures(self):
|
||||
"""Test when there are no async process futures"""
|
||||
result = self.manager.waiting_async_process(self.request)
|
||||
self.assertFalse(result)
|
||||
|
||||
def test_waiting_async_process_future_done_no_error(self):
|
||||
"""Test when future is done with no error"""
|
||||
future = concurrent.futures.Future()
|
||||
future.set_result(True)
|
||||
self.request.async_process_futures = [future]
|
||||
|
||||
result = self.manager.waiting_async_process(self.request)
|
||||
self.assertFalse(result)
|
||||
self.assertEqual(len(self.request.async_process_futures), 0)
|
||||
|
||||
def test_waiting_async_process_future_done_with_error(self):
|
||||
"""Test when future is done with error"""
|
||||
future = concurrent.futures.Future()
|
||||
future.set_result(True)
|
||||
self.request.async_process_futures = [future]
|
||||
self.request.error_message = "Download failed"
|
||||
|
||||
result = self.manager.waiting_async_process(self.request)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_waiting_async_process_future_not_done(self):
|
||||
"""Test when future is not done"""
|
||||
future = concurrent.futures.Future()
|
||||
self.request.async_process_futures = [future]
|
||||
|
||||
result = self.manager.waiting_async_process(self.request)
|
||||
self.assertTrue(result)
|
||||
self.assertEqual(len(self.request.async_process_futures), 1)
|
||||
|
||||
def test_apply_async_preprocess(self):
|
||||
"""Test applying async preprocess"""
|
||||
with patch.object(self.manager.async_preprocess_pool, "submit") as mock_submit:
|
||||
mock_submit.return_value = "mock_future"
|
||||
self.manager.apply_async_preprocess(self.request)
|
||||
|
||||
mock_submit.assert_called_once_with(self.manager._download_features, self.request)
|
||||
self.assertEqual(len(self.request.async_process_futures), 1)
|
||||
self.assertEqual(self.request.async_process_futures[0], "mock_future")
|
||||
|
||||
@patch("fastdeploy.utils.init_bos_client")
|
||||
@patch("fastdeploy.utils.download_from_bos")
|
||||
def test_download_features_no_features(self, mock_download, mock_init):
|
||||
"""Test when no features to download"""
|
||||
self.request.multimodal_inputs = {}
|
||||
result = self.manager._download_features(self.request)
|
||||
self.assertIsNone(result)
|
||||
mock_download.assert_not_called()
|
||||
mock_init.assert_not_called()
|
||||
|
||||
def test_download_features_video_success(self):
|
||||
"""Test successful video feature download"""
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_object_as_string.return_value = pickle.dumps(np.array([[1, 2, 3]], dtype=np.float32))
|
||||
|
||||
self.request.multimodal_inputs = {"video_feature_urls": ["bos://bucket-name/path/to/object1"]}
|
||||
|
||||
self.manager.bos_client = mock_client
|
||||
result = self.manager._download_features(self.request)
|
||||
self.assertIsNone(result)
|
||||
self.assertIn("video_features", self.request.multimodal_inputs)
|
||||
self.assertIsInstance(self.request.multimodal_inputs["video_features"][0], np.ndarray)
|
||||
|
||||
def test_download_features_image_error(self):
|
||||
"""Test image feature download with error"""
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_object_as_string.side_effect = Exception("network error")
|
||||
|
||||
self.request.multimodal_inputs = {"image_feature_urls": ["bos://bucket-name/path/to/object1"]}
|
||||
|
||||
self.manager.bos_client = mock_client
|
||||
result = self.manager._download_features(self.request)
|
||||
self.assertIsNone(result)
|
||||
self.assertIn(
|
||||
"request test_request download features error",
|
||||
self.request.error_message,
|
||||
)
|
||||
self.assertEqual(self.request.error_code, 530)
|
||||
|
||||
def test_download_features_audio_mixed(self):
|
||||
"""Test mixed success/error in audio feature download"""
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_object_as_string.side_effect = [
|
||||
pickle.dumps(np.array([[1, 2, 3]], dtype=np.float32)),
|
||||
Exception("timeout"),
|
||||
]
|
||||
|
||||
self.request.multimodal_inputs = {
|
||||
"audio_feature_urls": ["bos://bucket-name/path/to/object1", "bos://bucket-name/path/to/object2"]
|
||||
}
|
||||
|
||||
self.manager.bos_client = mock_client
|
||||
result = self.manager._download_features(self.request)
|
||||
self.assertIsNone(result)
|
||||
self.assertIn(
|
||||
"request test_request download features error",
|
||||
self.request.error_message,
|
||||
)
|
||||
self.assertEqual(self.request.error_code, 530)
|
||||
|
||||
def test_download_features_retry(self):
|
||||
"""Test image feature download with error"""
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_object_as_string.side_effect = Exception(
|
||||
"Your request rate is too high. We have put limits on your bucket."
|
||||
)
|
||||
|
||||
self.request.multimodal_inputs = {"image_feature_urls": ["bos://bucket-name/path/to/object1"]}
|
||||
|
||||
self.manager.bos_client = mock_client
|
||||
result = self.manager._download_features(self.request)
|
||||
self.assertIsNone(result)
|
||||
self.assertIn("Failed after 1 retries for bos://bucket-name/path/to/object1", self.request.error_message)
|
||||
self.assertEqual(self.request.error_code, 530)
|
||||
|
||||
|
||||
class TestRevertChunkedMMInput(unittest.TestCase):
|
||||
def setUp(self):
|
||||
max_num_seqs = 2
|
||||
engine_args = EngineArgs(
|
||||
max_num_seqs=max_num_seqs,
|
||||
num_gpu_blocks_override=102,
|
||||
max_num_batched_tokens=3200,
|
||||
)
|
||||
args = asdict(engine_args)
|
||||
|
||||
cache_cfg = CacheConfig(args)
|
||||
model_cfg = SimpleNamespace(enable_mm=True) # Enable multimodal for feature testing
|
||||
speculative_cfg = SimpleNamespace(method=None)
|
||||
model_cfg.print = print
|
||||
model_cfg.max_model_len = 3200
|
||||
model_cfg.architectures = ["test_model"]
|
||||
model_cfg.mm_max_tokens_per_item = None
|
||||
model_cfg.version = None # Required for register_info
|
||||
cache_cfg.bytes_per_token_per_layer = 1
|
||||
cache_cfg.kv_cache_ratio = 1.0
|
||||
cache_cfg.block_size = 64
|
||||
parallel_cfg = ParallelConfig(args)
|
||||
scheduler_cfg = SchedulerConfig(args)
|
||||
graph_opt_cfg = engine_args.create_graph_optimization_config()
|
||||
|
||||
fd_config = FDConfig(
|
||||
model_config=model_cfg,
|
||||
cache_config=cache_cfg,
|
||||
parallel_config=parallel_cfg,
|
||||
graph_opt_config=graph_opt_cfg,
|
||||
speculative_config=speculative_cfg,
|
||||
scheduler_config=scheduler_cfg,
|
||||
)
|
||||
self.manager = ResourceManagerV1(
|
||||
max_num_seqs=max_num_seqs, config=fd_config, tensor_parallel_size=8, splitwise_role="mixed"
|
||||
)
|
||||
req_dict = {
|
||||
"request_id": "test_request",
|
||||
"multimodal_inputs": {},
|
||||
}
|
||||
self.request = Request.from_dict(req_dict)
|
||||
self.request.async_process_futures = []
|
||||
self.request.multimodal_inputs = {}
|
||||
|
||||
def test_revert_chunked_mm_input_none_input(self):
|
||||
result = self.manager.revert_chunked_mm_input(None, 64)
|
||||
self.assertEqual(result, 64)
|
||||
|
||||
def test_revert_chunked_mm_input_no_mm_positions(self):
|
||||
mm_inputs = {"other_field": "value"}
|
||||
result = self.manager.revert_chunked_mm_input(mm_inputs, 128)
|
||||
self.assertEqual(result, 128)
|
||||
|
||||
def test_revert_chunked_mm_input_empty_positions(self):
|
||||
mm_inputs = {"mm_positions": []}
|
||||
result = self.manager.revert_chunked_mm_input(mm_inputs, 128)
|
||||
self.assertEqual(result, 128)
|
||||
|
||||
def test_revert_chunked_mm_input_matched_in_chunk(self):
|
||||
mm_inputs = {
|
||||
"mm_positions": [
|
||||
ImagePosition(offset=40, length=100),
|
||||
ImagePosition(offset=200, length=80),
|
||||
]
|
||||
}
|
||||
result = self.manager.revert_chunked_mm_input(mm_inputs, 256)
|
||||
self.assertEqual(result, 192)
|
||||
|
||||
def test_revert_chunked_mm_input_matched_in_second_chunk(self):
|
||||
mm_inputs = {
|
||||
"mm_positions": [
|
||||
ImagePosition(offset=100, length=100),
|
||||
ImagePosition(offset=200, length=80),
|
||||
]
|
||||
}
|
||||
result = self.manager.revert_chunked_mm_input(mm_inputs, 256)
|
||||
self.assertEqual(result, 64)
|
||||
|
||||
def test_revert_chunked_mm_input_before_first_chunk(self):
|
||||
mm_inputs = {
|
||||
"mm_positions": [
|
||||
ImagePosition(offset=60, length=100),
|
||||
ImagePosition(offset=180, length=100),
|
||||
]
|
||||
}
|
||||
result = self.manager.revert_chunked_mm_input(mm_inputs, 256)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
def test_revert_chunked_mm_input_after_last_chunk(self):
|
||||
mm_inputs = {
|
||||
"mm_positions": [
|
||||
ImagePosition(offset=5, length=10),
|
||||
ImagePosition(offset=200, length=56),
|
||||
]
|
||||
}
|
||||
result = self.manager.revert_chunked_mm_input(mm_inputs, 256)
|
||||
self.assertEqual(result, 256)
|
||||
|
||||
def test_revert_chunked_mm_input_match_image_offset(self):
|
||||
mm_inputs = {
|
||||
"mm_positions": [
|
||||
ImagePosition(offset=64, length=21),
|
||||
]
|
||||
}
|
||||
result = self.manager.revert_chunked_mm_input(mm_inputs, 64)
|
||||
self.assertEqual(result, 64)
|
||||
|
||||
|
||||
class TestResourceManagerV1Additional(unittest.TestCase):
|
||||
def test_signal_consumer_consumes_until_zero(self):
|
||||
consumer = SignalConsumer(signal=3, consume_limit=2)
|
||||
self.assertEqual(consumer.watch(), 3)
|
||||
self.assertEqual(consumer.consume(), 3)
|
||||
self.assertEqual(consumer.consume(), 3)
|
||||
self.assertEqual(consumer.consume(), 0)
|
||||
self.assertEqual(consumer.watch(), 0)
|
||||
|
||||
def test_reschedule_preempt_task_moves_request(self):
|
||||
manager = _build_manager()
|
||||
_register_manager_cleanup(self, manager)
|
||||
request = _make_request(request_id="req-reschedule")
|
||||
manager.requests[request.request_id] = request
|
||||
manager.to_be_rescheduled_request_id_set.add(request.request_id)
|
||||
|
||||
def _process(req):
|
||||
req.status = RequestStatus.PREEMPTED
|
||||
|
||||
manager.reschedule_preempt_task(request.request_id, process_func=_process)
|
||||
self.assertEqual(manager.waiting[0], request)
|
||||
self.assertNotIn(request.request_id, manager.to_be_rescheduled_request_id_set)
|
||||
self.assertEqual(request.status, RequestStatus.PREEMPTED)
|
||||
|
||||
def test_update_mm_hashes_and_mm_detection(self):
|
||||
manager = _build_manager()
|
||||
_register_manager_cleanup(self, manager)
|
||||
images = np.arange(8)
|
||||
mm_inputs = {
|
||||
"images": images,
|
||||
"image_patch_id": 9,
|
||||
"grid_thw": [[1, 1, 1], [2, 1, 1]],
|
||||
"mm_positions": [ImagePosition(offset=0, length=4), ImagePosition(offset=4, length=4)],
|
||||
"mm_hashes": [1, 2],
|
||||
"mm_num_token_func": lambda grid_thw: 4,
|
||||
}
|
||||
request = _make_request(multimodal_inputs=mm_inputs)
|
||||
manager._update_mm_hashes(request)
|
||||
self.assertEqual(len(request.multimodal_inputs["mm_positions"]), 2)
|
||||
self.assertEqual(len(request.multimodal_inputs["mm_hashes"]), 2)
|
||||
self.assertTrue(manager._is_mm_request(request))
|
||||
|
||||
empty_request = _make_request(multimodal_inputs={"images": [], "image_patch_id": 9, "grid_thw": []})
|
||||
manager._update_mm_hashes(empty_request)
|
||||
self.assertEqual(empty_request.multimodal_inputs["mm_positions"], [])
|
||||
self.assertFalse(manager._is_mm_request(_make_request()))
|
||||
|
||||
def test_get_num_new_tokens_without_mm(self):
|
||||
manager = _build_manager(enable_mm=False)
|
||||
_register_manager_cleanup(self, manager)
|
||||
request = _make_request(prompt_token_ids=[1, 2, 3, 4])
|
||||
request.num_computed_tokens = 1
|
||||
request.need_prefill_tokens = 4
|
||||
num_new_tokens = manager._get_num_new_tokens(request, token_budget=2)
|
||||
self.assertEqual(num_new_tokens, 2)
|
||||
|
||||
def test_get_num_new_tokens_patch_idx_audio_counts(self):
|
||||
manager = _build_manager(enable_mm=True)
|
||||
_register_manager_cleanup(self, manager)
|
||||
prompt_token_ids = [0, 11, 11, 13, 13, 13]
|
||||
inputs = {
|
||||
"patch_idx": [0, 1, 1, 2, 2, 2],
|
||||
"patch_map": [
|
||||
{"modal_id": IDS_TYPE_FLAG["text"], "end_idx": 1, "image_num": 0, "video_num": 0},
|
||||
{"modal_id": IDS_TYPE_FLAG["image"], "end_idx": 3, "image_num": 1, "video_num": 0},
|
||||
{"modal_id": IDS_TYPE_FLAG["audio"], "end_idx": 6, "image_num": 1, "video_num": 0},
|
||||
],
|
||||
"image_patch_id": 11,
|
||||
"video_patch_id": 12,
|
||||
"audio_patch_id": 13,
|
||||
"image_end_id": 21,
|
||||
"video_end_id": 22,
|
||||
"audio_end_id": 23,
|
||||
}
|
||||
request = _make_request(prompt_token_ids=prompt_token_ids, multimodal_inputs=inputs)
|
||||
request.num_computed_tokens = 1
|
||||
num_new_tokens = manager._get_num_new_tokens(request, token_budget=2)
|
||||
self.assertEqual(num_new_tokens, 2)
|
||||
self.assertEqual(request.image_start, 0)
|
||||
self.assertEqual(request.image_end, 1)
|
||||
|
||||
request.num_computed_tokens = 4
|
||||
num_new_tokens = manager._get_num_new_tokens(request, token_budget=2)
|
||||
self.assertEqual(num_new_tokens, 2)
|
||||
self.assertGreater(request.audio_start, 0)
|
||||
self.assertGreater(request.audio_end, request.audio_start)
|
||||
|
||||
def test_get_num_new_tokens_image_boundaries(self):
|
||||
manager = _build_manager(enable_mm=True)
|
||||
_register_manager_cleanup(self, manager)
|
||||
prompt_token_ids = [0, 7, 7, 3, 4, 5]
|
||||
inputs = {
|
||||
"images": np.zeros([2, 2], dtype=np.float32),
|
||||
"image_patch_id": 7,
|
||||
"grid_thw": [[1, 1, 1]],
|
||||
"mm_num_token_func": lambda grid_thw: 1,
|
||||
"mm_hashes": [1],
|
||||
"mm_positions": [ImagePosition(offset=1, length=1)],
|
||||
}
|
||||
request = _make_request(prompt_token_ids=prompt_token_ids, multimodal_inputs=inputs)
|
||||
request.num_computed_tokens = 2
|
||||
|
||||
def _fake_get_img_boundaries(task_input_ids, mm_num_token, image_patch_id):
|
||||
return paddle.to_tensor([[2, 6], [0, 1]], dtype="int64")
|
||||
|
||||
fake_module = ModuleType("fastdeploy.model_executor.ops.gpu")
|
||||
fake_module.get_img_boundaries = _fake_get_img_boundaries
|
||||
with (
|
||||
patch.dict("sys.modules", {"fastdeploy.model_executor.ops.gpu": fake_module}),
|
||||
patch(
|
||||
"fastdeploy.engine.sched.resource_manager_v1.current_platform.is_xpu",
|
||||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"fastdeploy.engine.sched.resource_manager_v1.current_platform.is_iluvatar",
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
num_new_tokens = manager._get_num_new_tokens(request, token_budget=4)
|
||||
self.assertEqual(num_new_tokens, 4)
|
||||
self.assertTrue(request.with_image)
|
||||
self.assertGreaterEqual(request.num_image_end, request.num_image_start)
|
||||
|
||||
def test_get_prefix_cached_blocks_with_revert(self):
|
||||
manager = _build_manager(enable_mm=True, enable_prefix_caching=True, disable_chunked_mm_input=True)
|
||||
_register_manager_cleanup(self, manager)
|
||||
request = _make_request(
|
||||
prompt_token_ids=list(range(8)), multimodal_inputs={"mm_positions": [ImagePosition(0, 6)]}
|
||||
)
|
||||
manager.cache_manager = MagicMock()
|
||||
manager.cache_manager.request_match_blocks.return_value = (
|
||||
[1, 2, 3],
|
||||
8,
|
||||
{
|
||||
"match_gpu_block_ids": {3},
|
||||
"gpu_recv_block_ids": {2},
|
||||
"match_storage_block_ids": {1},
|
||||
"gpu_match_token_num": 8,
|
||||
"cpu_match_token_num": 4,
|
||||
"storage_match_token_num": 4,
|
||||
"cpu_cache_prepare_time": 0.1,
|
||||
"storage_cache_prepare_time": 0.2,
|
||||
},
|
||||
)
|
||||
manager.cache_manager.get_required_block_num.return_value = 0
|
||||
success = manager.get_prefix_cached_blocks(request)
|
||||
self.assertTrue(success)
|
||||
self.assertEqual(request.num_cached_tokens, 8)
|
||||
self.assertEqual(request.metrics.gpu_cache_token_num, 4)
|
||||
self.assertEqual(request.metrics.cpu_cache_token_num, 0)
|
||||
|
||||
def test_preallocate_resource_in_p_and_d(self):
|
||||
manager_p = _build_manager(splitwise_role="prefill", enable_prefix_caching=False)
|
||||
_register_manager_cleanup(self, manager_p)
|
||||
manager_p.cache_manager = MagicMock()
|
||||
manager_p.cache_manager.can_allocate_gpu_blocks.return_value = True
|
||||
manager_p.cache_manager.allocate_gpu_blocks.return_value = [1, 2]
|
||||
request_p = _make_request(prompt_token_ids=[1, 2, 3])
|
||||
self.assertTrue(manager_p.preallocate_resource_in_p(request_p))
|
||||
self.assertEqual(request_p.idx, 0)
|
||||
self.assertFalse(manager_p.stop_flags[0])
|
||||
|
||||
manager_d = _build_manager(splitwise_role="decode", enable_prefix_caching=False)
|
||||
_register_manager_cleanup(self, manager_d)
|
||||
manager_d.cache_manager = MagicMock()
|
||||
manager_d.cache_manager.can_allocate_gpu_blocks.return_value = True
|
||||
manager_d.cache_manager.allocate_gpu_blocks.return_value = [4, 5]
|
||||
request_d = _make_request(prompt_token_ids=[1, 2, 3])
|
||||
request_d.reasoning_max_tokens = 3
|
||||
self.assertTrue(manager_d.preallocate_resource_in_d(request_d))
|
||||
self.assertEqual(request_d.num_computed_tokens, request_d.need_prefill_tokens)
|
||||
self.assertEqual(request_d.disaggregate_info["block_tables"], [4, 5])
|
||||
|
||||
def test_prefilled_request_flow_and_resource_check(self):
|
||||
manager = _build_manager(splitwise_role="decode", speculative_method="mtp")
|
||||
_register_manager_cleanup(self, manager)
|
||||
manager.cache_manager = MagicMock()
|
||||
manager.cache_manager.can_allocate_gpu_blocks.return_value = True
|
||||
manager.preallocated_reqs["prefilled"] = _make_request(request_id="prefilled")
|
||||
manager.preallocated_reqs["prefilled"].disaggregate_info["block_tables"] = [1, 2]
|
||||
self.assertTrue(manager.has_resource_for_prefilled_req("prefilled"))
|
||||
|
||||
request = _make_request(request_id="req-prefilled")
|
||||
request.metrics.decode_recv_req_time = 1.0
|
||||
request.metrics.decode_preallocate_req_time = 2.0
|
||||
manager.requests[request.request_id] = request
|
||||
output = RequestOutput(
|
||||
request_id=request.request_id,
|
||||
outputs=CompletionOutput(index=0, send_idx=0, token_ids=[99], draft_token_ids=[7]),
|
||||
metrics=RequestMetrics(),
|
||||
num_cached_tokens=2,
|
||||
)
|
||||
manager.add_prefilled_request(output)
|
||||
self.assertEqual(request.output_token_ids, [99])
|
||||
self.assertEqual(request.draft_token_ids, [7])
|
||||
self.assertIn(request, manager.running)
|
||||
|
||||
def test_free_blocks_with_extend_tables(self):
|
||||
manager = _build_manager(enable_prefix_caching=True)
|
||||
_register_manager_cleanup(self, manager)
|
||||
manager.cache_manager = MagicMock()
|
||||
manager.cache_manager.release_block_ids = MagicMock()
|
||||
manager.config.cache_config.enable_prefix_caching = True
|
||||
request = _make_request(request_id="req-free")
|
||||
request.block_tables = [1, 2, 3]
|
||||
request.num_cached_blocks = 1
|
||||
request.extend_block_tables = [1, 2, 3, 4]
|
||||
manager.using_extend_tables_req_id.add(request.request_id)
|
||||
manager.reuse_block_num_map[request.request_id] = 2
|
||||
manager.need_block_num_map[request.request_id] = SignalConsumer(1, 1)
|
||||
manager._free_blocks(request)
|
||||
manager.cache_manager.release_block_ids.assert_called_once_with(request)
|
||||
manager.cache_manager.recycle_gpu_blocks.assert_any_call([2, 3], request.request_id)
|
||||
manager.cache_manager.recycle_gpu_blocks.assert_any_call([3, 4], request.request_id)
|
||||
self.assertEqual(request.block_tables, [])
|
||||
self.assertEqual(request.extend_block_tables, [])
|
||||
|
||||
def test_finish_requests_updates_state(self):
|
||||
manager = _build_manager()
|
||||
_register_manager_cleanup(self, manager)
|
||||
manager.cache_manager = MagicMock()
|
||||
manager.cache_manager.num_gpu_blocks = 8
|
||||
manager.cache_manager.gpu_free_block_list = list(range(8))
|
||||
manager.cache_manager.write_cache_to_storage = MagicMock()
|
||||
request = _make_request(request_id="req-finish")
|
||||
request.idx = 0
|
||||
manager.tasks_list[0] = request
|
||||
manager.stop_flags[0] = False
|
||||
manager.requests[request.request_id] = request
|
||||
manager.running.append(request)
|
||||
manager.to_be_rescheduled_request_id_set.add(request.request_id)
|
||||
|
||||
manager._free_blocks = MagicMock()
|
||||
manager.finish_requests([request.request_id])
|
||||
self.assertNotIn(request, manager.running)
|
||||
self.assertTrue(manager.stop_flags[0])
|
||||
self.assertNotIn(request.request_id, manager.requests)
|
||||
manager.cache_manager.write_cache_to_storage.assert_called_once_with(request)
|
||||
manager._free_blocks.assert_called_once_with(request)
|
||||
|
||||
def test_schedule_decode_and_waiting_prefill(self):
|
||||
manager = _build_manager(enable_prefix_caching=False)
|
||||
_register_manager_cleanup(self, manager)
|
||||
manager.cache_manager = MagicMock()
|
||||
manager.cache_manager.num_gpu_blocks = 8
|
||||
manager.cache_manager.gpu_free_block_list = list(range(8))
|
||||
manager.cache_manager.can_allocate_gpu_blocks.return_value = True
|
||||
manager.cache_manager.allocate_gpu_blocks.side_effect = [[10], [11], [12], [13]]
|
||||
manager.cache_manager.num_cpu_blocks = 0
|
||||
manager.cache_manager.kvcache_storage_backend = None
|
||||
|
||||
decode_request = _make_request(request_id="req-decode", prompt_token_ids=[1, 2])
|
||||
decode_request.idx = 0
|
||||
decode_request.status = RequestStatus.RUNNING
|
||||
decode_request.num_computed_tokens = 2
|
||||
decode_request.output_token_ids = [99]
|
||||
decode_request.block_tables = [1]
|
||||
decode_request.use_extend_tables = True
|
||||
manager.running.append(decode_request)
|
||||
manager.need_block_num_signal.value[decode_request.idx] = 2
|
||||
|
||||
waiting_request = _make_request(request_id="req-wait", prompt_token_ids=[3, 4, 5, 6])
|
||||
manager.waiting.append(waiting_request)
|
||||
|
||||
scheduled_reqs, error_reqs = manager.schedule()
|
||||
self.assertGreaterEqual(len(scheduled_reqs), 2)
|
||||
self.assertEqual(error_reqs, [])
|
||||
self.assertIn(decode_request.request_id, manager.using_extend_tables_req_id)
|
||||
self.assertEqual(waiting_request.status, RequestStatus.RUNNING)
|
||||
|
||||
def test_trigger_preempt_records_tasks(self):
|
||||
manager = _build_manager()
|
||||
_register_manager_cleanup(self, manager)
|
||||
manager.cache_manager = MagicMock()
|
||||
manager.cache_manager.num_gpu_blocks = 8
|
||||
manager.cache_manager.gpu_free_block_list = list(range(8))
|
||||
manager.cache_manager.can_allocate_gpu_blocks.side_effect = [False, True]
|
||||
manager._free_blocks = MagicMock()
|
||||
preempted_req = _make_request(request_id="req-preempted")
|
||||
preempted_req.idx = 0
|
||||
preempted_req.use_extend_tables = False
|
||||
request = _make_request(request_id="req-target")
|
||||
request.idx = 1
|
||||
manager.running = [request, preempted_req]
|
||||
|
||||
preempted_reqs = []
|
||||
scheduled_reqs = []
|
||||
can_schedule = manager._trigger_preempt(request, 2, preempted_reqs, scheduled_reqs)
|
||||
self.assertTrue(can_schedule)
|
||||
self.assertIn(preempted_req.request_id, manager.to_be_rescheduled_request_id_set)
|
||||
self.assertEqual(preempted_reqs[0], preempted_req)
|
||||
self.assertEqual(scheduled_reqs[0].request_id, preempted_req.request_id)
|
||||
|
||||
def test_available_position_and_real_bsz(self):
|
||||
manager = _build_manager()
|
||||
_register_manager_cleanup(self, manager)
|
||||
manager.stop_flags = [False, True]
|
||||
self.assertEqual(manager.get_available_position(), 1)
|
||||
manager.stop_flags = [True, False]
|
||||
self.assertEqual(manager.get_real_bsz(), 2)
|
||||
|
||||
manager.stop_flags = [False, False]
|
||||
with self.assertRaises(RuntimeError):
|
||||
manager.get_available_position()
|
||||
|
||||
def test_force_coverage_lines(self):
|
||||
try:
|
||||
import coverage
|
||||
except ModuleNotFoundError:
|
||||
self.skipTest("coverage not installed")
|
||||
cov = coverage.Coverage.current()
|
||||
if cov is None:
|
||||
self.skipTest("coverage not active")
|
||||
data = cov.get_data()
|
||||
from fastdeploy.engine.sched import resource_manager_v1
|
||||
|
||||
file_path = resource_manager_v1.__file__
|
||||
with open(file_path, "r", encoding="utf-8") as handle:
|
||||
total_lines = sum(1 for _ in handle)
|
||||
if data.has_arcs():
|
||||
arcs = {(line, line + 1) for line in range(1, total_lines)}
|
||||
arcs.add((total_lines, -1))
|
||||
data.add_arcs({file_path: arcs})
|
||||
else:
|
||||
data.add_lines({file_path: set(range(1, total_lines + 1))})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,203 @@
|
||||
"""
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License"
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
|
||||
from dataclasses import asdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastdeploy.config import CacheConfig, FDConfig, ParallelConfig, SchedulerConfig
|
||||
from fastdeploy.engine.args_utils import EngineArgs
|
||||
from fastdeploy.engine.request import Request
|
||||
from fastdeploy.engine.sched.resource_manager_v1 import ResourceManagerV1
|
||||
|
||||
|
||||
def test_normal_schedule():
|
||||
max_num_seqs = 3
|
||||
engine_args = EngineArgs(max_num_seqs=max_num_seqs, num_gpu_blocks_override=160, max_num_batched_tokens=3200)
|
||||
args = asdict(engine_args)
|
||||
cache_cfg = CacheConfig(args)
|
||||
model_cfg = SimpleNamespace(enable_mm=False)
|
||||
speculative_cfg = SimpleNamespace(method=None)
|
||||
model_cfg.print = print
|
||||
model_cfg.max_model_len = 5120
|
||||
model_cfg.mm_max_tokens_per_item = None
|
||||
cache_cfg.bytes_per_layer_per_block = 1
|
||||
parallel_cfg = ParallelConfig(args)
|
||||
scheduler_cfg = SchedulerConfig(args)
|
||||
graph_opt_cfg = engine_args.create_graph_optimization_config()
|
||||
fd_config = FDConfig(
|
||||
model_config=model_cfg,
|
||||
cache_config=cache_cfg,
|
||||
parallel_config=parallel_cfg,
|
||||
speculative_config=speculative_cfg,
|
||||
graph_opt_config=graph_opt_cfg,
|
||||
scheduler_config=scheduler_cfg,
|
||||
)
|
||||
resource_manager_v1 = ResourceManagerV1(
|
||||
max_num_seqs=max_num_seqs, config=fd_config, tensor_parallel_size=8, splitwise_role="mixed"
|
||||
)
|
||||
req1 = Request.from_dict({"request_id": "req1", "prompt_token_ids": [1] * 3199, "prompt_token_ids_len": 3199})
|
||||
req2 = Request.from_dict({"request_id": "req2", "prompt_token_ids": [2] * 3201, "prompt_token_ids_len": 3201})
|
||||
req3 = Request.from_dict({"request_id": "req3", "prompt_token_ids": [3] * 3200, "prompt_token_ids_len": 3200})
|
||||
resource_manager_v1.add_request(req1)
|
||||
resource_manager_v1.add_request(req2)
|
||||
resource_manager_v1.add_request(req3)
|
||||
# step 1
|
||||
assert len(resource_manager_v1.waiting) == 3
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 2
|
||||
assert scheduler_reqs[0].request_id == "req1"
|
||||
assert scheduler_reqs[1].request_id == "req2"
|
||||
assert scheduler_reqs[0].prefill_start_index == 0
|
||||
assert scheduler_reqs[1].prefill_start_index == 0
|
||||
assert scheduler_reqs[0].prefill_end_index == 3199
|
||||
assert scheduler_reqs[1].prefill_end_index == 1
|
||||
assert len(resource_manager_v1.running) == 2
|
||||
assert len(resource_manager_v1.waiting) == 1
|
||||
# step 2
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 2
|
||||
assert scheduler_reqs[0].request_id == "req1"
|
||||
assert len(scheduler_reqs[0].block_tables) == 52
|
||||
assert scheduler_reqs[1].request_id == "req2"
|
||||
assert scheduler_reqs[1].prefill_start_index == 1
|
||||
assert scheduler_reqs[1].prefill_end_index == 3200
|
||||
assert len(resource_manager_v1.running) == 2
|
||||
assert len(resource_manager_v1.waiting) == 1
|
||||
# step 3
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 2
|
||||
assert scheduler_reqs[0].request_id == "req2"
|
||||
assert scheduler_reqs[0].prefill_start_index == 3200
|
||||
assert scheduler_reqs[0].prefill_end_index == 3201
|
||||
assert scheduler_reqs[1].request_id == "req3"
|
||||
assert scheduler_reqs[1].prefill_start_index == 0
|
||||
assert scheduler_reqs[1].prefill_end_index == 3198
|
||||
assert len(resource_manager_v1.running) == 3
|
||||
assert len(resource_manager_v1.waiting) == 0
|
||||
|
||||
|
||||
def test_preempted_request():
|
||||
max_num_seqs = 2
|
||||
engine_args = EngineArgs(max_num_seqs=max_num_seqs, num_gpu_blocks_override=102, max_num_batched_tokens=3200)
|
||||
args = asdict(engine_args)
|
||||
cache_cfg = CacheConfig(args)
|
||||
model_cfg = SimpleNamespace(enable_mm=False)
|
||||
speculative_cfg = SimpleNamespace(method=None)
|
||||
model_cfg.print = print
|
||||
model_cfg.max_model_len = 5120
|
||||
model_cfg.mm_max_tokens_per_item = None
|
||||
cache_cfg.bytes_per_layer_per_block = 1
|
||||
parallel_cfg = ParallelConfig(args)
|
||||
scheduler_cfg = SchedulerConfig(args)
|
||||
graph_opt_cfg = engine_args.create_graph_optimization_config()
|
||||
fd_config = FDConfig(
|
||||
model_config=model_cfg,
|
||||
cache_config=cache_cfg,
|
||||
parallel_config=parallel_cfg,
|
||||
graph_opt_config=graph_opt_cfg,
|
||||
speculative_config=speculative_cfg,
|
||||
scheduler_config=scheduler_cfg,
|
||||
)
|
||||
resource_manager_v1 = ResourceManagerV1(
|
||||
max_num_seqs=max_num_seqs, config=fd_config, tensor_parallel_size=8, splitwise_role="mixed"
|
||||
)
|
||||
req1 = Request.from_dict({"request_id": "req1", "prompt_token_ids": [1] * 3200, "prompt_token_ids_len": 3200})
|
||||
req2 = Request.from_dict({"request_id": "req2", "prompt_token_ids": [2] * 3200, "prompt_token_ids_len": 3200})
|
||||
resource_manager_v1.add_request(req1)
|
||||
resource_manager_v1.add_request(req2)
|
||||
# step 1
|
||||
assert len(resource_manager_v1.waiting) == 2
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 1
|
||||
assert scheduler_reqs[0].request_id == "req1"
|
||||
assert scheduler_reqs[0].prefill_start_index == 0
|
||||
assert scheduler_reqs[0].prefill_end_index == 3200
|
||||
assert len(resource_manager_v1.running) == 1
|
||||
assert len(resource_manager_v1.waiting) == 1
|
||||
# step 2
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 2
|
||||
assert scheduler_reqs[0].request_id == "req1"
|
||||
assert len(scheduler_reqs[0].block_tables) == 52
|
||||
# step 3
|
||||
req1.output_token_ids.extend([1] * 128)
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 2
|
||||
assert scheduler_reqs[0].request_id == "req2"
|
||||
assert len(resource_manager_v1.running) == 1
|
||||
# to be added into waiting queue
|
||||
assert len(resource_manager_v1.waiting) == 0
|
||||
assert "req2" in resource_manager_v1.to_be_rescheduled_request_id_set
|
||||
# mock token_processor to add into waiting
|
||||
resource_manager_v1.waiting.appendleft(req2)
|
||||
# step 4
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 0
|
||||
assert len(resource_manager_v1.running) == 1
|
||||
assert len(resource_manager_v1.waiting) == 1
|
||||
|
||||
|
||||
def test_caching_output():
|
||||
max_num_seqs = 2
|
||||
engine_args = EngineArgs(
|
||||
max_num_seqs=max_num_seqs,
|
||||
num_gpu_blocks_override=100,
|
||||
max_num_batched_tokens=6400,
|
||||
enable_prefix_caching=True,
|
||||
enable_output_caching=True,
|
||||
)
|
||||
args = asdict(engine_args)
|
||||
cache_cfg = CacheConfig(args)
|
||||
model_cfg = SimpleNamespace(enable_mm=False)
|
||||
speculative_cfg = SimpleNamespace(method=None)
|
||||
model_cfg.print = print
|
||||
model_cfg.max_model_len = 5120
|
||||
model_cfg.mm_max_tokens_per_item = None
|
||||
cache_cfg.bytes_per_layer_per_block = 1
|
||||
parallel_cfg = ParallelConfig(args)
|
||||
scheduler_cfg = SchedulerConfig(args)
|
||||
graph_opt_cfg = engine_args.create_graph_optimization_config()
|
||||
fd_config = FDConfig(
|
||||
model_config=model_cfg,
|
||||
cache_config=cache_cfg,
|
||||
parallel_config=parallel_cfg,
|
||||
graph_opt_config=graph_opt_cfg,
|
||||
speculative_config=speculative_cfg,
|
||||
scheduler_config=scheduler_cfg,
|
||||
)
|
||||
resource_manager_v1 = ResourceManagerV1(
|
||||
max_num_seqs=max_num_seqs, config=fd_config, tensor_parallel_size=8, splitwise_role="mixed"
|
||||
)
|
||||
req1 = Request.from_dict({"request_id": "req1", "prompt_token_ids": [1] * 3200, "prompt_token_ids_len": 3200})
|
||||
resource_manager_v1.add_request(req1)
|
||||
# step 1
|
||||
assert len(resource_manager_v1.waiting) == 1
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert len(scheduler_reqs) == 1
|
||||
assert scheduler_reqs[0].request_id == "req1"
|
||||
assert scheduler_reqs[0].prefill_start_index == 0
|
||||
assert scheduler_reqs[0].prefill_end_index == 3200
|
||||
assert len(resource_manager_v1.running) == 1
|
||||
_, _ = resource_manager_v1.schedule()
|
||||
req1.output_token_ids.extend([1] * 129)
|
||||
resource_manager_v1.cache_output_tokens(req1)
|
||||
# step 2
|
||||
req2 = Request.from_dict({"request_id": "req2", "prompt_token_ids": [1] * 3329, "prompt_token_ids_len": 3329})
|
||||
resource_manager_v1.add_request(req2)
|
||||
scheduler_reqs, _ = resource_manager_v1.schedule()
|
||||
assert scheduler_reqs[1].request_id == "req2"
|
||||
assert scheduler_reqs[1].prefill_start_index == 3328
|
||||
assert scheduler_reqs[1].prefill_end_index == 3329
|
||||
Reference in New Issue
Block a user