local-llm-server/llm_server/routes/queue.py

import pickle
import time
from typing import Tuple
from uuid import uuid4

import ujson as json
from redis import Redis

from llm_server import opts
from llm_server.cluster.cluster_config import cluster_config
from llm_server.custom_redis import RedisCustom, redis
from llm_server.database.database import get_token_ratelimit


def increment_ip_count(client_ip: str, redis_key):
    redis.hincrby(redis_key, client_ip, 1)


def decrement_ip_count(client_ip: str, redis_key):
    new_count = redis.hincrby(redis_key, client_ip, -1)
    if new_count <= 0:
        redis.hdel(redis_key, client_ip)


class RedisPriorityQueue:
    """
    A queue for a specific backend.
    """

    def __init__(self, name, db: int = 12):
        self.name = name
        self.redis = RedisCustom(name, db=db)

    def put(self, item, priority: int, selected_model: str, do_stream: bool = False):
        # TODO: remove this when we're sure nothing strange is happening
        assert item is not None
        assert priority is not None
        assert selected_model is not None

        event = DataEvent()

        # Check if the IP is already in the dictionary and if it has reached the limit
        ip_count = self.get_ip_request_count(item[1])
        _, simultaneous_ip = get_token_ratelimit(item[2])
        if ip_count and int(ip_count) >= simultaneous_ip and priority != 0:
            print(f'Rejecting request from {item[1]} - {ip_count} request queued.')
            return None  # reject the request

        timestamp = time.time()
        self.redis.zadd('queue', {json.dumps((item, event.event_id, selected_model, timestamp, do_stream)): -priority})
        return event

    def get(self):
        while True:
            data = self.redis.zpopmin('queue')
            if data:
                item = json.loads(data[0][0])
                return item
            time.sleep(0.1)  # wait for something to be added to the queue

    def __len__(self):
        return self.redis.zcard('queue')

    def get_ip_request_count(self, client_ip: str):
        """
        Get the number of requests in the queue from a specific IP.
        This is a bit inefficient since we iterate over the entire queue, but
        keeps the queue as a single point of truth instead of tracking a separate hashed
        set which can get confusing.
        If we run into slowdowns in the future, we should go back to the hashed set approach.
        :param client_ip:
        :return:
        """
        start_time = time.time()
        items = self.redis.zrange('queue', 0, -1)
        count = 0
        for item in items:
            item_data = json.loads(item)
            if item_data[0][1] == client_ip:
                count += 1
        elapsed_time = time.time() - start_time
        if elapsed_time > 0.5:
            raise Exception(f"!!! get_ip_request_count took {elapsed_time} seconds to execute !!!")
        return count

    def flush(self):
        self.redis.flush()

    def items(self):
        return self.redis.zrange('queue', 0, -1)

    def cleanup(self):
        now = time.time()
        for item in self.items():
            item_data = json.loads(item)
            timestamp = item_data[-2]
            if now - timestamp > opts.backend_generate_request_timeout:
                self.redis.zrem('queue', 0, item)
                event_id = item_data[1]
                event = DataEvent(event_id)
                event.set((False, None, 'closed'))
                print('Removed timed-out item from queue:', event_id)


class DataEvent:
    """
    Class to simplify pub/sub communication between consumers and producers (MASTERS and SLAVES lololololol).
    """

    def __init__(self, event_id=None):
        self.event_id = event_id if event_id else str(uuid4())
        self.redis = Redis(host='localhost', port=6379, db=14)
        self.pubsub = self.redis.pubsub()
        self.pubsub.subscribe(self.event_id)

    def set(self, data):
        self.redis.publish(self.event_id, pickle.dumps(data))

    def wait(self):
        for item in self.pubsub.listen():
            if item['type'] == 'message':
                return pickle.loads(item['data'])


def update_active_workers(key: str, operation: str):
    if operation == 'incr':
        redis.incr(f'active_gen_workers:{key}')
    elif operation == 'decr':
        redis.decr(f'active_gen_workers:{key}')
        if redis.get(f'active_gen_workers:{key}', default=0, dtype=int) < 0:
            redis.set(f'active_gen_workers:{key}', 0)


def incr_active_workers(selected_model: str, backend_url: str):
    update_active_workers(selected_model, 'incr')
    update_active_workers(backend_url, 'incr')


def decr_active_workers(selected_model: str, backend_url: str):
    update_active_workers(selected_model, 'decr')
    update_active_workers(backend_url, 'decr')


class PriorityQueue:
    """
    Helper class to wrangler all the different queues.
    """

    def __init__(self, backends: set = None):
        """
        Only have to load the backends once.
        :param backends:
        """
        self.redis = Redis(host='localhost', port=6379, db=9)
        if backends:
            for item in backends:
                self.redis.sadd('backends', item)

    def get_backends(self):
        return {x.decode('utf-8') for x in self.redis.smembers('backends')}

    def get_queued_ip_count(self, client_ip: str):
        count = 0
        for backend_url in self.get_backends():
            queue = RedisPriorityQueue(backend_url)
            count += queue.get_ip_request_count(client_ip)
        return count

    def put(self, backend_url, item: Tuple[dict, str, str, dict], priority: int, selected_model: str, do_stream: bool = False):
        queue = RedisPriorityQueue(backend_url)
        return queue.put(item, priority, selected_model, do_stream)

    def activity(self):
        lines = []
        status_redis = RedisCustom('worker_status')
        for worker in status_redis.keys():
            lines.append((worker, status_redis.getp(worker)))
        return sorted(lines)

    def len(self, model_name):
        count = 0
        backends_with_models = set()
        for k in self.get_backends():
            info = cluster_config.get_backend(k)
            if info.get('model') == model_name:
                backends_with_models.add(k)
        for backend_url in backends_with_models:
            count += len(RedisPriorityQueue(backend_url))
        return count

    def __len__(self):
        count = 0
        p = set()
        for backend_url in self.get_backends():
            queue = RedisPriorityQueue(backend_url)
            p.add((backend_url, len(queue)))
            count += len(queue)
        return count

    def flush(self):
        for k in self.redis.keys():
            q = json.loads(self.redis.get(k))
            q.flush()
            self.redis.set(k, json.dumps(q))

    def flush_db(self):
        self.redis.flushdb()


priority_queue = PriorityQueue()