FastDeploy/fastdeploy/engine/sched/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 threading
import time
from collections import deque
from collections.abc import Iterable
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Union

import numpy as np
import paddle

from fastdeploy.engine.request import Request, RequestStatus, RequestType
from fastdeploy.engine.resource_manager import ResourceManager
from fastdeploy.utils import llm_logger


@dataclass
class ScheduledDecodeTask:
    """
    Task for allocating new blocks to decode.
    """

    idx: int
    request_id: str
    block_tables: list[int]
    task_type: RequestType = RequestType.DECODE


@dataclass
class ScheduledPreemptTask:
    """
    Task for terminating inference to recycle resource.
    """

    idx: int
    request_id: str
    task_type: RequestType = RequestType.PREEMPTED


class ResourceManagerV1(ResourceManager):
    """
    Resource manager for scheduler v1.
    In scheduler v1, all gpu blocks are managed by PrefixCacheManager.
    Tasks sent to worker are divided into 3 types, PREFILL、DECODE and PREEMPTED.
    For prefill task, the worker infer with one step and then stopped for this query if not all prompt tokens are computed.
    For decode task, the work continues to decode until allocated blocks are exhausted.
    For preempted task, the work reset all inputs to terminate the inference.
    """

    def __init__(self, max_num_seqs, config, tensor_parallel_size, splitwise_role, local_data_parallel_id=0):
        super(ResourceManagerV1, self).__init__(
            max_num_seqs, config, tensor_parallel_size, splitwise_role, local_data_parallel_id
        )
        # req_id -> Request
        self.config = config
        self.requests: dict[str, Request] = {}
        # Priority queues for requests.
        self.waiting: deque[Request] = deque()
        self.running: list[Request] = []
        self.finish_execution_pool = ThreadPoolExecutor(max_workers=1)
        self.lock = threading.Lock()
        self.to_be_rescheduled_request_id_set = set()

    def allocated_slots(self, request: Request):
        return len(request.block_tables) * self.config.cache_config.block_size

    def get_new_block_nums(self, request: Request, num_new_tokens: int):
        self.check_and_free_block_tables()
        return (
            request.num_computed_tokens + num_new_tokens + self.config.cache_config.block_size - 1
        ) // self.config.cache_config.block_size - len(request.block_tables)

    def _prepare_prefill_task(self, request, new_token_num):
        request.prefill_start_index = request.num_computed_tokens
        request.prefill_end_index = request.num_computed_tokens + new_token_num
        request.task_type = RequestType.PREFILL
        return request

    def _prepare_decode_task(self, request):
        return ScheduledDecodeTask(idx=request.idx, request_id=request.request_id, block_tables=request.block_tables)

    def _prepare_preempt_task(self, request):
        return ScheduledPreemptTask(idx=request.idx, request_id=request.request_id)

    def reschedule_preempt_task(self, request_id):
        with self.lock:
            if request_id in self.to_be_rescheduled_request_id_set and request_id in self.requests:
                request = self.requests[request_id]
                self.waiting.appendleft(request)
                self.to_be_rescheduled_request_id_set.remove(request_id)

    def _trigger_preempt(self, request, num_new_blocks, preempted_reqs, scheduled_reqs):
        can_schedule = True
        while True:
            if not self.cache_manager.can_allocate_gpu_blocks(num_new_blocks):
                preempted_req = self.running.pop()
                preempted_req.status = RequestStatus.PREEMPTED
                preempted_req.num_computed_tokens = 0
                preempted_req.prefill_block_num = 0
                self._free_blocks(preempted_req)
                self.to_be_rescheduled_request_id_set.add(preempted_req.request_id)
                preempted_reqs.append(preempted_req)
                scheduled_reqs.append(self._prepare_preempt_task(preempted_req))
                if preempted_req == request:
                    # No more request to preempt.
                    can_schedule = False
                    break
            else:
                # The request can be scheduled.
                can_schedule = True
                break
        return can_schedule

    def _get_num_new_tokens(self, request, token_budget):
        num_new_tokens = request.need_prefill_tokens - request.num_computed_tokens
        num_new_tokens = min(num_new_tokens, token_budget)

        if not self.config.enable_mm:
            return num_new_tokens

        inputs = request.multimodal_inputs
        request.with_image = False
        # Compatible with scenarios without images and videos.
        if inputs["images"] is None:
            return num_new_tokens

        input_ids_lst = request.prompt_token_ids + request.output_token_ids
        input_ids = paddle.to_tensor(input_ids_lst, dtype="int64")
        input_ids = paddle.to_tensor(input_ids_lst, dtype="int64")
        image_patch_id = inputs["image_patch_id"]

        if request.multimodal_img_boundaries is None:
            grid_thw = []
            for one in inputs["grid_thw"]:
                if one[0] == 1:
                    grid_thw.append(one)
                else:
                    grid_thw.extend([[2, one[1], one[2]]] * (one[0] // 2))

            grid_thw = paddle.to_tensor(grid_thw, dtype="int64")
            from fastdeploy.model_executor.ops.gpu import get_img_boundaries

            request.multimodal_img_boundaries = get_img_boundaries(
                task_input_ids=input_ids, grid_thw=grid_thw, image_patch_id=image_patch_id
            ).numpy()

            grid_thw = grid_thw.numpy().reshape([-1, 3])
            inputs["grid_thw"] = grid_thw

        grid_thw = inputs["grid_thw"]
        img_boundaries_idx = request.multimodal_img_boundaries[0]
        img_num_per_boundary = request.multimodal_img_boundaries[1]
        ori_prompt_len = img_boundaries_idx[-1].item()
        pre_end_idx = request.num_computed_tokens
        new_end_idx = pre_end_idx + num_new_tokens
        if new_end_idx < ori_prompt_len and input_ids[new_end_idx - 1] == image_patch_id:
            boundary_idx = np.searchsorted(img_boundaries_idx, new_end_idx, side="left").item()
            if boundary_idx == len(img_boundaries_idx):
                new_end_idx = ori_prompt_len
            else:
                new_end_idx = img_boundaries_idx[boundary_idx].item()
        elif new_end_idx >= ori_prompt_len and paddle.sum(input_ids[pre_end_idx:new_end_idx] == image_patch_id):
            new_end_idx = ori_prompt_len
        num_new_tokens = new_end_idx - pre_end_idx

        image_mask = input_ids[pre_end_idx:new_end_idx] == image_patch_id
        request.with_image = image_mask.any()
        if request.with_image:
            pre_boundary_idx = np.searchsorted(img_boundaries_idx, pre_end_idx, side="left").item()
            if pre_boundary_idx == len(img_boundaries_idx):
                request.num_image_start = img_num_per_boundary[-1]
            else:
                pre_boundary_idx = (
                    pre_boundary_idx if pre_end_idx == img_boundaries_idx[pre_boundary_idx] else pre_boundary_idx - 1
                )
                request.num_image_start = img_num_per_boundary[pre_boundary_idx]

            new_boundary_idx = np.searchsorted(img_boundaries_idx, new_end_idx, side="left").item()
            if new_boundary_idx == len(img_boundaries_idx):
                request.num_image_end = img_num_per_boundary[-1]
            else:
                new_boundary_idx = (
                    new_boundary_idx if new_end_idx == img_boundaries_idx[new_boundary_idx] else new_boundary_idx - 1
                )
                request.num_image_end = img_num_per_boundary[new_boundary_idx]

            request.image_type_ids_start = np.sum(grid_thw[: request.num_image_start, 0])
            request.image_type_ids_end = np.sum(grid_thw[: request.num_image_end, 0])
            request.image_start = np.sum(np.prod(grid_thw[: request.num_image_start], axis=1))
            request.image_end = np.sum(np.prod(grid_thw[: request.num_image_end], axis=1))
        return num_new_tokens

    def exist_prefill(self, scheduled_reqs):
        for request in scheduled_reqs:
            if request.task_type == RequestType.PREFILL:
                return True
        return False

    def schedule(self):
        with self.lock:
            scheduled_reqs: list[Request] = []
            preempted_reqs: list[Request] = []
            token_budget = self.config.max_num_batched_tokens

            # First, schedule the RUNNING requests.
            req_index = 0
            num_decoding_req_nums = 0
            while req_index < len(self.running) and token_budget > 0:
                request = self.running[req_index]
                if request.num_computed_tokens >= request.need_prefill_tokens:  # to be decoding
                    if request.num_total_tokens > request.need_prefill_tokens:  # has generated tokens
                        request.num_computed_tokens = request.num_total_tokens - 1
                    else:  # prefill finished
                        if (
                            self.config.cache_config.enable_prefix_caching
                            and request.get("prefill_block_num", None) is None
                        ):
                            # update prefill cache blocks for prefix caching
                            request.prefill_block_num = len(request.block_tables)
                            self.cache_manager.update_cache_blocks(request, self.config.cache_config.block_size)
                    if (
                        self.allocated_slots(request) - request.num_total_tokens
                        <= self.config.cache_config.prealloc_dec_block_slot_num_threshold
                    ):
                        # Allocation for next decoding blocks
                        if self.cache_manager.can_allocate_gpu_blocks(self.config.cache_config.enc_dec_block_num):
                            llm_logger.debug(
                                f"schedule decoding task: {request} request.num_total_tokens {request.num_total_tokens} request.num_computed_tokens {request.num_computed_tokens}"
                            )
                            request.block_tables.extend(
                                self.cache_manager.allocate_gpu_blocks(self.config.cache_config.enc_dec_block_num)
                            )
                            # Prepare decoding task
                            scheduled_reqs.append(self._prepare_decode_task(request))
                        else:
                            # Not enough blocks to allocate, trigger preemption
                            can_schedule = self._trigger_preempt(
                                request, self.config.cache_config.enc_dec_block_num, preempted_reqs, scheduled_reqs
                            )
                            if not can_schedule:
                                break
                            # Allocation for next decoding blocks
                            request.block_tables.extend(
                                self.cache_manager.allocate_gpu_blocks(self.config.cache_config.enc_dec_block_num)
                            )
                            # Prepare decoding task
                            scheduled_reqs.append(self._prepare_decode_task(request))
                        num_decoding_req_nums += 1
                        token_budget -= 1
                else:  # need to prefill
                    llm_logger.debug(
                        f"scheduler prefill task: {request} request.need_prefill_tokens {request.need_prefill_tokens} request.num_computed_tokens {request.num_computed_tokens}"
                    )
                    num_new_tokens = self._get_num_new_tokens(request, token_budget)
                    num_new_block = self.get_new_block_nums(request, num_new_tokens)
                    # Allocate blocks to prefill
                    if self.cache_manager.can_allocate_gpu_blocks(num_new_block):
                        request.block_tables.extend(self.cache_manager.allocate_gpu_blocks(num_new_block))
                        # Prepare prefill task
                        scheduled_reqs.append(self._prepare_prefill_task(request, num_new_tokens))
                    else:
                        can_schedule = self._trigger_preempt(request, num_new_block, preempted_reqs, scheduled_reqs)
                        if not can_schedule:
                            break
                        request.block_tables.extend(self.cache_manager.allocate_gpu_blocks(num_new_block))
                        # Prepare prefill task
                        scheduled_reqs.append(self._prepare_prefill_task(request, num_new_tokens))
                    token_budget -= num_new_tokens
                    request.num_computed_tokens += num_new_tokens
                req_index += 1
            # schedule the WAITING requests.
            if not preempted_reqs:
                while self.waiting and token_budget > 0:
                    if len(self.running) == self.max_num_seqs:
                        break
                    if self.config.enable_mm and self.exist_prefill(scheduled_reqs):
                        break
                    request = self.waiting[0]
                    if request.status == RequestStatus.WAITING:
                        # Enable prefix caching
                        if self.config.cache_config.enable_prefix_caching:
                            success = self.get_prefix_cached_blocks(request)
                            if not success:
                                break

                        num_new_tokens = self._get_num_new_tokens(request, token_budget)
                        num_new_block = self.get_new_block_nums(request, num_new_tokens)
                        # Allocate blocks to prefill
                        if self.cache_manager.can_allocate_gpu_blocks(num_new_block):
                            if not request.get("skip_allocate", False):
                                request.block_tables.extend(self.cache_manager.allocate_gpu_blocks(num_new_block))
                            self.waiting.popleft()
                            self.running.append(request)
                            scheduled_reqs.append(self._prepare_prefill_task(request, num_new_tokens))
                            request.inference_start_time = time.time()
                            request.schedule_start_time = time.time()
                            token_budget -= num_new_tokens
                            request.num_computed_tokens += num_new_tokens
                            request.status = RequestStatus.RUNNING
                            allocated_position = self.get_available_position()
                            request.idx = allocated_position
                            self.tasks_list[allocated_position] = request
                            self.stop_flags[allocated_position] = False
                            self.req_dict[request.request_id] = allocated_position
                        else:
                            break
                    elif request.status == RequestStatus.PREEMPTED:
                        request.need_prefill_tokens = (
                            request.num_total_tokens
                        )  # Before preempted task rescheduled, preempted task has been sent to engine, no more tokens are output, here num_total_tokens should be static and correct
                        num_new_tokens = self._get_num_new_tokens(request, token_budget)
                        num_new_block = self.get_new_block_nums(request, num_new_tokens)
                        # Allocate blocks to prefill
                        if self.cache_manager.can_allocate_gpu_blocks(num_new_block):
                            request.block_tables.extend(self.cache_manager.allocate_gpu_blocks(num_new_block))
                            self.waiting.popleft()
                            self.running.append(request)
                            scheduled_reqs.append(self._prepare_prefill_task(request, num_new_tokens))
                            token_budget -= num_new_tokens
                            request.num_computed_tokens += num_new_tokens
                            request.status = RequestStatus.RUNNING
                        else:
                            break
                    else:
                        llm_logger.error("Unknown request status type")
            if scheduled_reqs:
                llm_logger.debug(f"schedued_reqs: {scheduled_reqs}")
            return scheduled_reqs

    def get_available_position(self) -> int:
        position = 0
        while position < self.max_num_seqs:
            if self.stop_flags[position] is True:
                return position
            position += 1
        raise RuntimeError("No available position is available for new request")

    def get_real_bsz(self) -> int:
        for i in range(self.max_num_seqs - 1, -1, -1):
            if not self.stop_flags[i]:
                self.real_bsz = i + 1
                break
        return self.real_bsz

    def get_prefix_cached_blocks(self, request: Request):
        """
        set prefix cached information for the given request
        """
        try:
            cache_prepare_time = time.time()
            (common_block_ids, matched_token_num, hit_info) = self.cache_manager.request_match_blocks(
                request, self.config.cache_config.block_size
            )

            matched_block_num = len(common_block_ids)
            no_cache_block_num = self.cache_manager.get_required_block_num(
                request.prompt_token_ids_len - matched_token_num,
                self.config.cache_config.block_size,
            )

            request.num_cached_tokens = matched_token_num
            request.gpu_cache_token_num = hit_info["gpu_cache_blocks"] * self.config.cache_config.block_size
            request.cpu_cache_token_num = hit_info["cpu_cache_blocks"] * self.config.cache_config.block_size
            request.cache_info = (matched_block_num, no_cache_block_num)
            request.block_tables = common_block_ids
            request.skip_allocate = False

            if matched_token_num == request.prompt_token_ids_len:
                request.num_computed_tokens = matched_token_num - 1
                request.skip_allocate = True
            else:
                request.num_computed_tokens = matched_token_num
            request.cache_prepare_time = time.time() - cache_prepare_time
            return True
        except Exception as e:
            llm_logger.error(f"prefix match blocks error: {e}, waiting reschedule...")
            return False

    def add_request(self, request: Request) -> None:
        with self.lock:
            self.waiting.append(request)
            self.requests[request.request_id] = request

    def _free_blocks(self, request: Request):
        if self.config.cache_config.enable_prefix_caching:
            # TODO(chengyanfu): support cache ouput blocks for prefix caching
            self.cache_manager.release_block_ids_async(request)
            self.cache_manager.recycle_gpu_blocks(request.block_tables[request.prefill_block_num :])
        else:
            self.cache_manager.recycle_gpu_blocks(request.block_tables)
        request.block_tables = []

    def finish_requests_async(self, request_ids: Union[str, Iterable[str]]):
        return self.finish_execution_pool.submit(self.finish_requests, request_ids)

    def finish_requests(self, request_ids: Union[str, Iterable[str]]):
        llm_logger.info(f"recycle resources for requests: {request_ids}")
        try:
            with self.lock:
                if isinstance(request_ids, str):
                    request_ids = (request_ids,)
                else:
                    request_ids = set(request_ids)
                for req_id in request_ids:
                    request = self.requests.get(req_id)
                    if request is None:
                        # Invalid request ID.
                        continue
                    if request in self.running:  # normally run and finished
                        self.running.remove(request)
                        request.status = RequestStatus.FINISHED
                        self._free_blocks(request)
                    if (
                        request.request_id in self.to_be_rescheduled_request_id_set
                    ):  # finished after preempted, blocks have been recycled.
                        self.to_be_rescheduled_request_id_set.remove(
                            request.request_id
                        )  # just remove from to_be_rescheduled_request_id_set
                    if (
                        request in self.waiting
                    ):  # after finished, this request still scheduled from preempted to waiting, unexpected error, should not be here
                        raise RuntimeError(f"request {request.request_id} scheduled into waiting list, after finished")
                    self.tasks_list[request.idx] = None
                    self.stop_flags[request.idx] = True
                    del self.requests[req_id]
        except Exception as e:
            llm_logger.error(e)