FastDeploy/fastdeploy/scheduler/global_scheduler.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.
"""


from typing import List, Optional, Dict, Tuple
import traceback
import threading
import time
from datetime import datetime
import random
import uuid
import crcmod
from redis import ConnectionPool
from fastdeploy.scheduler.storage import AdaptedRedis
from fastdeploy.engine.request import Request, RequestOutput
from fastdeploy.scheduler.data import ScheduledRequest, ScheduledResponse
from fastdeploy.scheduler.workers import Workers, Task
from fastdeploy.utils import llm_logger
from fastdeploy.scheduler import utils


class GlobalScheduler(object):
    """
    A distributed task scheduler that manages request/response queues using Redis.

    This class provides functionality for:
    - Enqueuing and dequeuing requests
    - Load balancing across multiple scheduler instances
    - Handling request/response lifecycle
    - Maintaining worker health checks
    """

    def __init__(self,
                 host: str,
                 port: int,
                 db: int,
                 password: Optional[str],
                 topic: str,
                 ttl: int,
                 min_load_score: float,
                 load_shrads_num: int,
                 enable_chunked_prefill: bool,
                 max_num_partial_prefills: int,
                 max_long_partial_prefills: int,
                 long_prefill_token_threshold: int,
                 ):
        """
        Initialize the GlobalScheduler with Redis connection and scheduling parameters.

        Args:
            host: Redis server hostname
            port: Redis server port
            db: Redis database number
            password: Optional password for Redis authentication
            topic: Base topic name for queue namespacing
            ttl: Time-to-live in seconds for Redis keys
            min_load_score: Minimum load score for task assignment
            load_shrads_num: Number of shards for load balancing table
            enable_chunked_prefill: Whether to enable chunked prefill processing
            max_num_partial_prefills: Maximum number of partial prefills allowed
            max_long_partial_prefills: Maximum number of long partial prefills allowed
            long_prefill_token_threshold: Token count threshold for long prefills

        Initializes:
            - Redis connection pool and client
            - Worker threads for request/response handling
            - Load balancing and request stealing mechanisms
            - Response tracking structures
        """

        self.topic = topic
        self.ttl = ttl
        self.min_load_score = min_load_score
        self.load_shrads_num = load_shrads_num

        self.enable_chunked_prefill = enable_chunked_prefill
        self.max_num_partial_prefills = max_num_partial_prefills
        self.max_long_partial_prefills = max_long_partial_prefills
        self.long_prefill_token_threshold = long_prefill_token_threshold

        self.blpop_request_timeout = 2
        self.blpop_response_timeout = 10

        self.crc16_mutex = threading.Lock()
        self.crc16 = crcmod.predefined.Crc('ccitt-false')
        self.load_slot_for_getting_request = 0
        self.load_start = 0   # const
        self.load_num = 50    # const

        connection_pool = ConnectionPool(
            host=host, port=port, db=db, password=password, max_connections=10)
        self.client = AdaptedRedis(connection_pool=connection_pool)

        self.name = self._generate_scheduler_name()
        self.keep_alive_workers = threading.Thread(
            target=self._keep_alive, daemon=True)
        self.keep_alive_workers.start()

        self.put_requests_workers = Workers(
            "put_requests_workers", self._put_requests_worker, 20)
        self.put_requests_workers.start(1)

        self.put_results_workers = Workers(
            "put_results_workers", self._put_results_worker, 300)
        self.put_results_workers.start(1)

        self.mutex = threading.Lock()
        self.local_response_not_empty = threading.Condition(self.mutex)
        self.local_responses: Dict[str, List[ScheduledResponse]] = dict()
        self.stolen_requests: Dict[str, ScheduledRequest] = dict()

        self.get_response_workers = threading.Thread(
            target=self._get_results_worker, daemon=True)
        self.get_response_workers.start()

        llm_logger.info(
            f"Scheduler: name={self.name} redis_version={self.client.version}")

    def _get_hash_slot(self, data: str) -> int:
        data = data.encode("utf-8")
        with self.crc16_mutex:
            self.crc16.update(data)
            crc_value = self.crc16.crcValue
            self.crc16.crcValue = self.crc16.initCrc
        return crc_value

    def _instance_name(self, scheduler_name: str) -> str:
        """
        Generate the Redis key name for a scheduler instance.

        Args:
            scheduler_name: Name of the scheduler instance

        Returns:
            Formatted Redis key name
        """
        return f"{self.topic}.ins.{scheduler_name}"

    def _generate_scheduler_name(self) -> str:
        """
        Generate a unique name for this scheduler instance.

        Uses hostname/IP and timestamp to create a unique identifier,
        then registers it in Redis with TTL.

        Returns:
            Unique scheduler name string
        """
        try:
            _, name = utils.get_hostname_ip()
        except Exception as e:
            llm_logger.warning(
                f"Scheduler encountered an error while resolving the IP address. {e}")
            name = str(uuid.uuid4())

        size = len(name)
        now = time.time()
        local_time = datetime.fromtimestamp(now)
        formatted_time = local_time.strftime(
            "%Y-%m-%d %H:%M:%S") + f"{local_time.microsecond // 1000:03d}"

        count = 1
        while True:
            if self.client.set(self._instance_name(name), formatted_time, ex=self.ttl, nx=True):
                break
            name = f"{name[:size]}:{count}"
            count += 1
        return name

    def _keep_alive(self):
        """
        Background thread that periodically updates the scheduler's TTL in Redis.

        Runs in a loop with interval of TTL/2 to maintain instance registration.
        """
        interval_time = self.ttl / 2
        while True:
            try:
                now = time.time()
                local_time = datetime.fromtimestamp(now)
                formatted_time = local_time.strftime(
                    "%Y-%m-%d %H:%M:%S") + f"{local_time.microsecond // 1000:03d}"
                self.client.set(self._instance_name(self.name),
                                formatted_time, ex=self.ttl)
            except Exception as e:
                llm_logger.error(f"Scheduler keep alive failed: {e}")
                interval_time = self.ttl / 10

            time.sleep(interval_time)
            interval_time = self.ttl / 2

    def _scheduler_name_from_request_queue(self, request_queue: str) -> str:
        """
        Extract scheduler name from a request queue name.

        Args:
            request_queue: Full request queue name

        Returns:
            The scheduler name portion of the queue name
        """
        prefix_len = len(f"{self.topic}.req.")
        return request_queue[prefix_len:]

    def _request_queue_name(self, scheduler_name: Optional[str] = None) -> str:
        """
        Generate the Redis request queue name for a scheduler.

        Args:
            scheduler_name: Optional specific scheduler name, defaults to current instance

        Returns:
            Formatted request queue name
        """
        if scheduler_name is None:
            return f"{self.topic}.req.{self.name}"
        return f"{self.topic}.req.{scheduler_name}"

    def _response_queue_name(self, scheduler_name: Optional[str] = None) -> str:
        """
        Generate the Redis response queue name for a scheduler.

        Args:
            scheduler_name: Optional specific scheduler name, defaults to current instance

        Returns:
            Formatted response queue name
        """
        if scheduler_name is None:
            return f"{self.topic}.resp.{self.name}"
        return f"{self.topic}.resp.{scheduler_name}"

    def _load_table_name(self, request_queue_name: Optional[str] = None, slot: Optional[int] = None) -> str:
        """
        Get the Redis sorted set name used for load balancing.

        Returns:
            The load score key name
        """
        if request_queue_name is None:
            request_queue_name = self._request_queue_name()

        if slot is None:
            slot = self._get_hash_slot(
                request_queue_name) % self.load_shrads_num
        else:
            slot %= self.load_shrads_num
        return f"{self.topic}.load.{slot}"

    @staticmethod
    def calc_required_blocks(token_num, block_size):
        """
        Calculate the number of blocks needed for a given number of tokens.

        Args:
            token_num: Number of tokens
            block_size: Size of each block

        Returns:
            Number of blocks required (rounded up)
        """
        return (token_num + block_size - 1) // block_size

    @staticmethod
    def _mark_request(request: ScheduledRequest):
        """
        Mark a stolen request with the original queue name.

        Args:
            request: The request to mark
        """
        request.request_id = f"mark<{request.request_queue_name}>{request.request_id}"

    @staticmethod
    def _unmark_response(response: ScheduledResponse, request_queue_name: str):
        """
        Remove marking from a response that came from a stolen request.

        Args:
            response: The response to unmark
            request_queue_name: Original request queue name
        """
        mark = f"mark<{request_queue_name}>"
        if not response.request_id.startswith(mark):
            return
        response.request_id = response.request_id[len(mark):]

    def _put_requests_worker(self, tasks: List[Task]) -> List[Task]:
        """
        Worker method that adds requests to the shared Redis cache.

        Args:
            tasks: List of tasks containing requests to enqueue

        Returns:
            List of processed tasks (some may be marked as duplicates)
        """
        duplicate = False
        requests: List[ScheduledRequest] = []
        with self.mutex:
            for task in tasks:
                request = ScheduledRequest(
                    task.raw, self._request_queue_name(), self._response_queue_name())
                task.raw = None

                if request.request_id in self.local_responses:
                    task.reason = "duplicate request_id"
                    duplicate = True
                    continue
                requests.append(request)
                self.local_responses[request.request_id] = []

        if len(requests) > 0:
            serialized_requests = [request.serialize() for request in requests]
            self.client.rpush(self._request_queue_name(), *
                              serialized_requests, ttl=self.ttl)
            self.client.zincrby(self._load_table_name(),
                                len(serialized_requests), self.name,
                                rem_amount=0, ttl=self.ttl)
            llm_logger.info(
                f"Scheduler has enqueued some requests: {requests}")

        if duplicate:
            llm_logger.warning(
                "Scheduler has received some duplicated requests: "
                f"{[task for task in tasks if task.reason is not None]}")
        return tasks

    def put_requests(self, requests: List[Request]) -> List[Tuple[str, Optional[str]]]:
        """
        Public method to add new requests to the scheduler.

        Args:
            requests: List of Request objects to schedule

        Returns:
            List of tuples containing (request_id, error_reason) for each request
        """
        tasks: List[Task] = []
        for request in requests:
            task = Task(request.request_id, request)
            tasks.append(task)

        self.put_requests_workers.add_tasks(tasks)
        results = self.put_requests_workers.get_results(10, 0.001)
        return [(result.id, result.reason) for result in results]

    def get_requests(self, available_blocks, block_size, reserved_output_blocks,
                     max_num_batched_tokens, batch=1) -> List[Request]:
        """
        Get requests from the shared cache based on available resources.

        Args:
            available_blocks: Number of available processing blocks
            block_size: Size of each processing block
            reserved_output_blocks: Blocks reserved for output
            max_num_batched_tokens: Maximum tokens that can be batched
            batch: Preferred batch size

        Returns:
            List of Request objects ready for processing
        """

        if available_blocks <= reserved_output_blocks or batch < 1:
            llm_logger.debug(
                f"Scheduler's resource are insufficient: available_blocks={available_blocks} "
                f"reserved_output_blocks={reserved_output_blocks} batch={batch} "
                f"max_num_batched_tokens={max_num_batched_tokens}")
            return []

        mini_batch = (batch + 1) // 2
        batches = []
        for _ in range(2):
            if batch >= mini_batch:
                batches.append(mini_batch)
                batch -= mini_batch
                continue

            if batch > 0:
                batches.append(batch)
                batch = 0

        local_request_queue_name = self._request_queue_name()
        serialized_requests: List[Tuple[str, bytes]] = []
        for bs in batches:
            elements = self.client.lpop(
                local_request_queue_name, bs, ttl=self.ttl)
            if elements is None:
                break
            self.client.zincrby(self._load_table_name(), -
                                len(elements), self.name, rem_amount=0, ttl=self.ttl)
            serialized_requests += [(local_request_queue_name, element)
                                    for element in elements]

        extend_scheduler_names = []
        if len(serialized_requests) == 0 and len(batches) > 0:
            for _ in range(min(5, self.load_shrads_num)):
                serialized_members = self.client.zrangebyscore(
                    self._load_table_name(
                        slot=self.load_slot_for_getting_request),
                    self.min_load_score,
                    float("+inf"),
                    start=self.load_start,
                    num=self.load_num)
                self.load_slot_for_getting_request += 1
                if len(serialized_members) > 0:
                    break

            members = [member.decode("utf-8") for member in serialized_members]
            if len(members) > 0:
                extend_scheduler_names = random.sample(
                    members, k=min(10, len(members)))
                extend_scheduler_names = [
                    name for name in extend_scheduler_names if name != self.name]

        # find lucky one
        if len(extend_scheduler_names) > 0:
            lucky = random.choice(extend_scheduler_names)
            lucky_request_queue_name = self._request_queue_name(lucky)

            elements = self.client.lpop(lucky_request_queue_name, batches[0])
            if elements is not None and len(elements) > 0:
                self.client.zincrby(
                    self._load_table_name(
                        request_queue_name=lucky_request_queue_name),
                    -len(elements), lucky, rem_amount=0, ttl=self.ttl)
                serialized_requests += [(lucky_request_queue_name, element)
                                        for element in elements]
                llm_logger.info(
                    f"Scheduler {self.name} has stolen some requests from another lucky one. "
                    f"(name={lucky} num={len(serialized_requests)})")
            else:
                exist_num = self.client.exists(self._instance_name(lucky))
                if exist_num == 0:
                    if self.client.zrem(
                            self._load_table_name(
                                request_queue_name=lucky_request_queue_name),
                            lucky):
                        llm_logger.info(
                            f"Scheduler {lucky} has been removed")

        # blocked read
        if len(serialized_requests) == 0:
            request_queue_names = [local_request_queue_name]
            request_queue_names += [
                self._request_queue_name(name) for name in extend_scheduler_names]

            element = self.client.blpop(
                request_queue_names, self.blpop_request_timeout)
            if element is None:
                return []
            request_queue_name = element[0].decode("utf-8")
            scheduler_name = self._scheduler_name_from_request_queue(
                request_queue_name)
            self.client.zincrby(
                self._load_table_name(request_queue_name=request_queue_name),
                -1, scheduler_name, rem_amount=0, ttl=self.ttl)
            serialized_requests.append((request_queue_name, element[1]))
            if scheduler_name != self.name:
                llm_logger.info(
                    f"Scheduler {self.name} has stolen a request from another scheduler. (name={scheduler_name})")

        long_partial_requests = 0
        short_partial_requests = 0
        required_total_blocks = 0
        current_prefill_tokens = 0
        remaining_request: List[Tuple[str, bytes]] = []
        scheduled_requests: List[ScheduledRequest] = []
        for request_queue_name, serialized_request in serialized_requests:
            if len(remaining_request) > 0:
                remaining_request.append(
                    (request_queue_name, serialized_request))
                continue

            request: ScheduledRequest = ScheduledRequest.unserialize(
                serialized_request)
            required_input_blocks = self.calc_required_blocks(
                request.prompt_tokens_ids_len, block_size)

            current_prefill_tokens += request.prompt_tokens_ids_len
            required_total_blocks += required_input_blocks + reserved_output_blocks

            if required_total_blocks > available_blocks:
                remaining_request.append(
                    (request_queue_name, serialized_request))
                continue

            if self.enable_chunked_prefill:
                if request.prompt_tokens_ids_len > self.long_prefill_token_threshold:
                    long_partial_requests += 1
                    if long_partial_requests > self.max_long_partial_prefills:
                        remaining_request.append(
                            (request_queue_name, serialized_request))
                        continue
                else:
                    short_partial_requests += 1

                if short_partial_requests + long_partial_requests > self.max_num_partial_prefills:
                    remaining_request.append(
                        (request_queue_name, serialized_request))
                    continue
            else:
                if current_prefill_tokens > max_num_batched_tokens:
                    remaining_request.append(
                        (request_queue_name, serialized_request))
                    continue

            scheduled_requests.append(request)

        if len(scheduled_requests) > 0:
            with self.mutex:
                for request in scheduled_requests:
                    if request.request_queue_name == local_request_queue_name:
                        continue

                    self._mark_request(request)
                    if request.request_id not in self.stolen_requests:
                        self.stolen_requests[request.request_id] = request
                        continue

                    llm_logger.error(
                        f"Scheduler has received a duplicate request from others: {request}")

        requests: List[Request] = [
            request.raw for request in scheduled_requests]
        if len(remaining_request) > 0:
            group: Dict[str, List] = dict()
            for request_queue_name, serialized_request in remaining_request:
                if request_queue_name not in group:
                    group[request_queue_name] = []
                group[request_queue_name].append(serialized_request)

            for request_queue_name, serialized_requests in group.items():
                self.client.lpush(request_queue_name, *
                                  serialized_requests)
                scheduler_name = self._scheduler_name_from_request_queue(
                    request_queue_name)
                self.client.zincrby(
                    self._load_table_name(
                        request_queue_name=request_queue_name),
                    len(serialized_requests), scheduler_name, ttl=self.ttl)

            llm_logger.info(
                f"Scheduler has put remaining request into the queue: {len(remaining_request)}")
            if len(requests) == 0:
                llm_logger.debug(
                    f"Scheduler has put all just-pulled request into the queue: {len(remaining_request)}")

        if len(requests) > 0:
            llm_logger.info(
                f"Scheduler has pulled some request: {[request.request_id for request in requests]}")
        return requests

    def _put_results_worker(self, tasks: List[Task]):
        """
        Worker method that adds task results to the appropriate queues.

        Args:
            tasks: List of completed tasks with results
        """
        # count = 0  # for test

        with self.mutex:
            local_request_ids = set(self.local_responses.keys())

            stolen_request_id_request_queue = dict()
            stolen_request_id_response_queue = dict()
            for request_id, request in self.stolen_requests.items():
                stolen_request_id_request_queue[request_id] = request.request_queue_name
                stolen_request_id_response_queue[request_id] = request.response_queue_name

        finished_request_ids: List[str] = list()
        local_responses: Dict[str, List[ScheduledResponse]] = dict()
        stolen_responses: Dict[str, List[bytes]] = dict()

        for task in tasks:
            response = ScheduledResponse(task.raw)
            if response.finished:
                finished_request_ids.append(response.request_id)

            if response.request_id in local_request_ids:
                if response.request_id not in local_responses:
                    local_responses[response.request_id] = []
                local_responses[response.request_id].append(response)
                continue

            if response.request_id in stolen_request_id_request_queue:
                response_queue_name = stolen_request_id_response_queue[response.request_id]
                request_queue_name = stolen_request_id_request_queue[response.request_id]
                self._unmark_response(response, request_queue_name)

                if response_queue_name not in stolen_responses:
                    stolen_responses[response_queue_name] = []
                stolen_responses[response_queue_name].append(
                    response.serialize())
                continue

            llm_logger.error(
                f"Scheduler has recieved a non-existent response from engine: {[response]}")

        with self.mutex:
            for request_id, responses in local_responses.items():
                self.local_responses[request_id] += responses
                # count += len(responses)  # for test

            for request_id in finished_request_ids:
                if request_id in self.stolen_requests:
                    del self.stolen_requests[request_id]

            if len(local_responses) > 0:
                self.local_response_not_empty.notify_all()

        if len(finished_request_ids) > 0:
            llm_logger.info(
                f"Scheduler has received some finished responses: {finished_request_ids}")

        for response_queue_name, responses in stolen_responses.items():
            self.client.rpush(response_queue_name, *responses, ttl=self.ttl)
            # count += len(responses)  # for test
        # return [Task("", count)]  # for test

    def put_results(self, results: List[RequestOutput]):
        """
        Public method to add processing results back to the scheduler.

        Args:
            results: List of RequestOutput objects to return
        """
        tasks: List[Task] = [Task(result.request_id, result)
                             for result in results]
        self.put_results_workers.add_tasks(tasks)

        # ---- for test ----
        # task_results = self.put_results_workers.get_results(10, 0.001)
        # amount = 0
        # for task_result in task_results:
        #     amount += task_result.raw
        # return amount
        # ---- for test ----

    def _get_results_worker(self):
        """
        Background worker that continuously fetches results from Redis.

        Handles both bulk and blocking operations for efficiency.
        Runs in an infinite loop until scheduler shutdown.
        """
        while True:
            try:
                serialized_responses = self.client.lpop(
                    self._response_queue_name(), 300, ttl=self.ttl)

                if serialized_responses is None or len(serialized_responses) == 0:
                    element = self.client.blpop(
                        [self._response_queue_name()], self.blpop_response_timeout)
                    if element is None or len(element) == 0:
                        continue
                    serialized_responses = [element[1]]

                responses: Dict[str, List[ScheduledResponse]] = dict()
                for serialized_response in serialized_responses:
                    response = ScheduledResponse.unserialize(
                        serialized_response)
                    if response.request_id not in responses:
                        responses[response.request_id] = []
                    responses[response.request_id].append(response)

                with self.mutex:
                    for request_id, contents in responses.items():
                        if request_id not in self.local_responses:
                            llm_logger.error(
                                "Scheduler has received some non-existent response from the queue. "
                                f"response:{contents} queue:{self._response_queue_name()}")
                            continue
                        self.local_responses[request_id] += contents
                    self.local_response_not_empty.notify_all()
            except Exception as e:
                llm_logger.error(f"Scheduler get_results_worker exception: {e} "
                                 f"traceback: {traceback.format_exc()}")

    def get_results(self) -> Dict[str, List[RequestOutput]]:
        """
        Retrieve all available results from the distributed scheduler.

        This method:
        - Waits for new responses using a condition variable (timeout=0.001s)
        - Returns all currently available responses
        - Automatically removes completed requests from local tracking
        - Logs finished requests

        Behavior Details:
        1. For first call with less than 64 pending responses, returns empty dict
        2. Subsequent calls return all available responses
        3. Uses thread-safe operations with condition variables
        4. Automatically cleans up completed request tracking

        Returns:
            Dict[str, List[RequestOutput]]: 
                A dictionary where:
                - Key is the request ID
                - Value is a list of RequestOutput objects for that request
                Completed requests are automatically removed from tracking

        Note:
            - Thread-safe operation using condition variables
            - Short timeout avoids blocking while maintaining responsiveness
            - First call may return empty to batch small responses
            - Automatically logs finished requests via llm_logger
        """
        first = True

        def _get_results() -> Dict[str, List[ScheduledResponse]]:
            nonlocal first
            responses: Dict[str, List[ScheduledResponse]] = dict()

            count = 0
            for _, contents in self.local_responses.items():
                count += len(contents)

            if first and count < 64:
                first = False
                return responses

            request_ids = list(self.local_responses.keys())
            for request_id in request_ids:
                responses[request_id] = self.local_responses[request_id]
                self.local_responses[request_id] = []
            return responses

        with self.local_response_not_empty:
            responses: Dict[str, List[ScheduledResponse]] = self.local_response_not_empty.wait_for(
                _get_results, 0.001)

            results: Dict[str, List[RequestOutput]] = dict()
            for request_id, resps in responses.items():
                finished = False
                results[request_id] = []
                for resp in resps:
                    results[request_id].append(resp.raw)
                    finished |= resp.finished

                if finished:
                    del self.local_responses[request_id]
                    llm_logger.info(
                        f"Scheduler has pulled a finished response: {[request_id]}")
            return results

    def reset(self):
        """
        Reset the scheduler to its initial state by:
        1. Clearing all Redis queues associated with this scheduler instance
        2. Removing this instance from the load balancing table
        3. Clearing in-memory tracking of responses and stolen requests

        This method is thread-safe and should be called when:
        - The scheduler needs to be cleanly restarted
        - Recovering from critical errors
        - Preparing for graceful shutdown

        Effects:
        - Deletes the request and response queues in Redis
        - Removes this scheduler's entry from the load balancing sorted set
        - Clears the local_responses dictionary tracking pending responses
        - Clears the stolen_requests dictionary tracking requests taken from other schedulers

        Note: 
        - Uses the scheduler's mutex to ensure thread safety
        - Does not affect other scheduler instances in the cluster
        - After reset, the scheduler will need to be reinitialized to be usable again
        """
        with self.mutex:
            self.client.delete(self._request_queue_name(),
                               self._response_queue_name())
            self.client.zrem(self._load_table_name(), self.name)
            self.local_responses = dict()
            self.stolen_requests = dict()
        llm_logger.info("Scheduler has been reset")