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

import threading
import time

from llm_server import opts
from llm_server.llm.generator import generator
from llm_server.routes.cache import redis
from llm_server.routes.queue import DataEvent, decr_active_workers, decrement_ip_count, incr_active_workers, increment_ip_count, priority_queue


def worker():
    while True:
        need_to_wait()
        (request_json_body, client_ip, token, parameters), event_id = priority_queue.get()
        need_to_wait()

        increment_ip_count(client_ip, 'processing_ips')
        incr_active_workers()

        if not request_json_body:
            # This was a dummy request from the websocket handler.
            # We're going to let the websocket handler decrement processing_ips and active_gen_workers.
            continue

        try:
            success, response, error_msg = generator(request_json_body)
            event = DataEvent(event_id)
            event.set((success, response, error_msg))
        finally:
            decrement_ip_count(client_ip, 'processing_ips')
            decr_active_workers()


def start_workers(num_workers: int):
    i = 0
    for _ in range(num_workers):
        t = threading.Thread(target=worker)
        t.daemon = True
        t.start()
        i += 1
    print(f'Started {i} inference workers.')


def need_to_wait():
    # We need to check the number of active workers since the streaming endpoint may be doing something.
    active_workers = redis.get('active_gen_workers', int, 0)
    s = time.time()
    while active_workers >= opts.concurrent_gens:
        time.sleep(0.01)
    e = time.time()
    if e - s > 0.5:
        print(f'Worker was delayed {e - s} seconds.')