local-llm-server/llm_server/workers/inferencer.py

import threading
import time
from uuid import uuid4

from llm_server.cluster.cluster_config import cluster_config
from llm_server.custom_redis import redis, RedisCustom
from llm_server.llm.generator import generator
from llm_server.routes.queue import DataEvent, decr_active_workers, decrement_ip_count, incr_active_workers, increment_ip_count, RedisPriorityQueue, PriorityQueue, priority_queue


def worker(backend_url):
    queue = RedisPriorityQueue(backend_url)
    while True:
        (request_json_body, client_ip, token, parameters), event_id, selected_model = queue.get()
        backend_info = cluster_config.get_backend(backend_url)
        if not selected_model:
            selected_model = backend_info['model']

        increment_ip_count(client_ip, 'processing_ips')
        incr_active_workers(selected_model, backend_url)

        try:
            if not request_json_body:
                # This was a dummy request from the streaming handlers.
                # The worker will let the handler do the streaming instead
                # of the worker. The worker will block until the handler
                # is finished. Since a lot of ratelimiting and stats are
                # based off the number of active workers, we must keep
                # the generation based off the workers.
                pubsub = redis.pubsub()
                pubsub.subscribe(event_id)
                redis.publish(event_id, 'begin')
                for item in pubsub.listen():
                    if item['type'] == 'message' and item['data'].decode('utf-8') == 'finished':
                        break
                    time.sleep(0.1)
            else:
                # Normal inference (not streaming).
                success, response, error_msg = generator(request_json_body, backend_url)
                event = DataEvent(event_id)
                event.set((success, response, error_msg))
        finally:
            decrement_ip_count(client_ip, 'processing_ips')
            decr_active_workers(selected_model, backend_url)


def start_workers(cluster: dict):
    i = 0
    for item in cluster:
        for _ in range(item['concurrent_gens']):
            t = threading.Thread(target=worker, args=(item['backend_url'],))
            t.daemon = True
            t.start()
            i += 1
    print(f'Started {i} inference workers.')