58 lines
2.4 KiB
Python
58 lines
2.4 KiB
Python
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)
|
|
|
|
print('Worker starting processing for', client_ip)
|
|
|
|
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.')
|