FastDeploy/tests/v1/test_resource_manager_v1.py

"""
# 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
        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 = []
        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()