apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2026-04-23 00:17:25 +08:00
Code Issues Actions 19 Packages Projects Releases Wiki Activity
Files
develop
FastDeploy/tests/v1/test_resource_manager_v1.py
T
kevin b19754c5f4 [Scheduler][BugFix] Fix token_budget calculation to use actual decode request count (#7499) 
* [Scheduler][BugFix] Fix token_budget calculation to use actual decode request count

## Motivation

当前 `token_budget` 的计算方式存在两个问题：

1. **预扣过多**：budget 按 `max_num_seqs * tokens_per_seq` 预扣，而不是 running 队列中实际处于 decode 阶段的请求数，导致 prefill 可用的 token 数被低估。
2. **循环内重复扣减**：decode 分支固定执行 `token_budget -= 1`，在 spec decode 场景下（`tokens_per_seq > 1`）每个 decode 请求只扣 1，少扣了 `num_speculative_tokens` 个；此外，当 running 队列中 prefill 请求耗尽 budget 后，排在其后的 decode 请求会被循环退出条件 `token_budget > 0` 提前跳过，导致调度漏发。

## Modifications

- `resource_manager_v1.py`
  - 新增 `_is_decoding(request)` 内部方法，封装 `num_computed_tokens >= need_prefill_tokens` 判断，全文统一使用
  - 调度前统计 running 队列中真实的 decode 请求数 `num_running_decode_reqs`，以 `num_running_decode_reqs * tokens_per_seq` 一次性预扣 budget，替代原来的 `max_num_seqs * tokens_per_seq`
  - 去掉 decode 分支内的 `token_budget -= 1`（已在循环前整体预扣）
  - 修改循环退出条件：decode 请求不受 `token_budget <= 0` 限制，仅 prefill 请求在 budget 耗尽时退出

- `config.py`
  - 修复 `max_num_batched_tokens` 的合法性校验，考虑 spec decode 场景下 `tokens_per_seq = num_speculative_tokens + 1`，改为检查 `max_num_batched_tokens >= max_num_seqs * tokens_per_seq`

## Usage or Command

```bash
# 普通启动（非spec decode，行为不变）
python -m fastdeploy.entrypoints.openai.api_server \
  --max-num-batched-tokens 8192 \
  --max-num-seqs 256 \
  ...

# spec decode 场景（tokens_per_seq = num_speculative_tokens + 1）
# 确保 max_num_batched_tokens >= max_num_seqs * tokens_per_seq，否则启动报错
python -m fastdeploy.entrypoints.openai.api_server \
  --max-num-batched-tokens 8192 \
  --max-num-seqs 256 \
  --num-speculative-tokens 4 \
  ...
```

* [FDConfig][BugFix] Fix AttributeError when speculative_config is SimpleNamespace without num_speculative_tokens

## Motivation

当测试中使用 `SimpleNamespace(method=None)` 构造 `speculative_config` 时，
`config.py` 的 `check()` 方法直接访问 `self.speculative_config.num_speculative_tokens`，
导致 `AttributeError: 'types.SimpleNamespace' object has no attribute 'num_speculative_tokens'`。
影响以下测试文件：
- tests/v1/test_resource_manager_v1.py
- tests/eplb/test_eplb_utils.py
- tests/eplb/test_experts_manager.py
- tests/v1/cache_manager/test_prefix_cache.py
- tests/v1/test_schedule_output.py

## Modifications

- `fastdeploy/config.py`: 使用 `getattr(..., "num_speculative_tokens", 0)` 兜底，
  防止 speculative_config 对象缺少该属性时崩溃
- 测试文件：将 `speculative_config=SimpleNamespace(method=None)` 统一改为
  `speculative_config=None`，与无投机解码场景语义一致

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 17:43:28 +08:00

729 lines
30 KiB
Python
Raw Permalink Blame History

"""
# 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, "enable_compat"):
paddle.enable_compat = lambda scope=None: None
from fastdeploy.config import CacheConfig, FDConfig, ParallelConfig, SchedulerConfig
from fastdeploy.engine.args_utils import EngineArgs
from fastdeploy.engine.request import (
BatchRequest,
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
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=None,
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 = []
batch_request = BatchRequest()
can_schedule = manager._trigger_preempt(request, 2, preempted_reqs, batch_request)
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(batch_request.requests[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()
Reference in New Issue View Git Blame Copy Permalink