cyberes
/
local-llm-server
Archived
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity

Merge cluster to master #3

Merged
cyberes merged 163 commits from cluster into master 2023-10-27 19:19:22 -06:00
Conversation 0 Commits 163 Files Changed 91 +195 -157
10 changed files with 195 additions and 157 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit e8964fcfd2 - Show all commits

13
daemon.py
View File

@ -3,13 +3,12 @@ import sys
import time
from pathlib import Path
from redis import Redis
from llm_server.cluster.cluster_config import cluster_config
from llm_server.cluster.redis_cycle import redis_cycler_db
from llm_server.cluster.stores import redis_running_models
from llm_server.config.load import load_config, parse_backends
from llm_server.custom_redis import redis
from llm_server.database.create import create_db
from llm_server.routes.queue import priority_queue
from llm_server.routes.v1.generate_stats import generate_stats
from llm_server.workers.threader import start_background
@ -21,11 +20,8 @@ else:
config_path = Path(script_path, 'config', 'config.yml')
if __name__ == "__main__":
flushed_keys = redis.flush()
print('Flushed', len(flushed_keys), 'keys from Redis.')
redis_cycler_db.flushall()
redis_running_models.flush()
Redis().flushall()
print('Flushed Redis.')
success, config, msg = load_config(config_path)
if not success:
@ -34,7 +30,6 @@ if __name__ == "__main__":
create_db()
priority_queue.flush()
cluster_config.clear()
cluster_config.load(parse_backends(config))

4
llm_server/config/load.py
View File

@ -3,11 +3,13 @@ import sys
import openai
import llm_server
from llm_server import opts
from llm_server.config.config import ConfigLoader, config_default_vars, config_required_vars
from llm_server.custom_redis import redis
from llm_server.database.conn import database
from llm_server.database.database import get_number_of_rows
from llm_server.routes.queue import PriorityQueue
def load_config(config_path):
@ -54,6 +56,8 @@ def load_config(config_path):
for item in config['cluster']:
opts.cluster_workers += item['concurrent_gens']
llm_server.routes.queue.priority_queue = PriorityQueue([x['backend_url'] for x in config['cluster']])
if opts.openai_expose_our_model and not opts.openai_api_key:
print('If you set openai_epose_our_model to false, you must set your OpenAI key in openai_api_key.')
sys.exit(1)

111
llm_server/routes/openai/chat_completions.py
View File

@ -74,7 +74,7 @@ def openai_chat_completions():
event = None
if not handler.is_client_ratelimited():
# Add a dummy event to the queue and wait for it to reach a worker
event = priority_queue.put((None, handler.client_ip, handler.token, None, handler.backend_url), handler.token_priority, handler.selected_model)
event = priority_queue.put(handler.backend_url, (None, handler.client_ip, handler.token, None), handler.token_priority, handler.selected_model)
if not event:
log_to_db(
handler.client_ip,
@ -107,63 +107,64 @@ def openai_chat_completions():
oai_string = generate_oai_string(30)
def generate():
response = generator(msg_to_backend, handler.backend_url)
generated_text = ''
partial_response = b''
for chunk in response.iter_content(chunk_size=1):
partial_response += chunk
if partial_response.endswith(b'\x00'):
json_strs = partial_response.split(b'\x00')
for json_str in json_strs:
if json_str:
try:
json_obj = json.loads(json_str.decode())
new = json_obj['text'][0].split(handler.prompt + generated_text)[1]
generated_text = generated_text + new
except IndexError:
# ????
continue
try:
response = generator(msg_to_backend, handler.backend_url)
generated_text = ''
partial_response = b''
for chunk in response.iter_content(chunk_size=1):
partial_response += chunk
if partial_response.endswith(b'\x00'):
json_strs = partial_response.split(b'\x00')
for json_str in json_strs:
if json_str:
try:
json_obj = json.loads(json_str.decode())
new = json_obj['text'][0].split(handler.prompt + generated_text)[1]
generated_text = generated_text + new
except IndexError:
# ????
continue
data = {
"id": f"chatcmpl-{oai_string}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"delta": {
"content": new
},
"finish_reason": None
}
]
}
yield f'data: {json.dumps(data)}\n\n'
yield 'data: [DONE]\n\n'
end_time = time.time()
elapsed_time = end_time - start_time
log_to_db(
handler.client_ip,
handler.token,
handler.prompt,
generated_text,
elapsed_time,
handler.parameters,
r_headers,
response_status_code,
r_url,
handler.backend_url,
)
data = {
"id": f"chatcmpl-{oai_string}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"delta": {
"content": new
},
"finish_reason": None
}
]
}
yield f'data: {json.dumps(data)}\n\n'
yield 'data: [DONE]\n\n'
end_time = time.time()
elapsed_time = end_time - start_time
log_to_db(
handler.client_ip,
handler.token,
handler.prompt,
generated_text,
elapsed_time,
handler.parameters,
r_headers,
response_status_code,
r_url,
handler.backend_url,
)
finally:
# After completing inference, we need to tell the worker we
# are finished.
if event_id: # may be None if ratelimited.
redis.publish(event_id, 'finished')
else:
print('event_id was None!')
return Response(generate(), mimetype='text/event-stream')
except Exception:
traceback.print_exc()
return 'INTERNAL SERVER', 500
finally:
# After completing inference, we need to tell the worker we
# are finished.
if event_id: # may be None if ratelimited.
redis.publish(event_id, 'finished')
else:
print('event_id was None!')

105
llm_server/routes/openai/completions.py
View File

@ -102,7 +102,7 @@ def openai_completions():
event = None
if not handler.is_client_ratelimited():
# Add a dummy event to the queue and wait for it to reach a worker
event = priority_queue.put((None, handler.client_ip, handler.token, None, handler.backend_url), handler.token_priority, handler.selected_model)
event = priority_queue.put(handler.backend_url, (None, handler.client_ip, handler.token, None), handler.token_priority, handler.selected_model)
if not event:
log_to_db(
handler.client_ip,
@ -135,61 +135,62 @@ def openai_completions():
oai_string = generate_oai_string(30)
def generate():
generated_text = ''
partial_response = b''
for chunk in response.iter_content(chunk_size=1):
partial_response += chunk
if partial_response.endswith(b'\x00'):
json_strs = partial_response.split(b'\x00')
for json_str in json_strs:
if json_str:
try:
json_obj = json.loads(json_str.decode())
new = json_obj['text'][0].split(handler.prompt + generated_text)[1]
generated_text = generated_text + new
except IndexError:
# ????
continue
try:
generated_text = ''
partial_response = b''
for chunk in response.iter_content(chunk_size=1):
partial_response += chunk
if partial_response.endswith(b'\x00'):
json_strs = partial_response.split(b'\x00')
for json_str in json_strs:
if json_str:
try:
json_obj = json.loads(json_str.decode())
new = json_obj['text'][0].split(handler.prompt + generated_text)[1]
generated_text = generated_text + new
except IndexError:
# ????
continue
data = {
"id": f"cmpl-{oai_string}",
"object": "text_completion",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"delta": {
"content": new
},
"finish_reason": None
}
]
}
yield f'data: {json.dumps(data)}\n\n'
yield 'data: [DONE]\n\n'
end_time = time.time()
elapsed_time = end_time - start_time
data = {
"id": f"cmpl-{oai_string}",
"object": "text_completion",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"delta": {
"content": new
},
"finish_reason": None
}
]
}
yield f'data: {json.dumps(data)}\n\n'
yield 'data: [DONE]\n\n'
end_time = time.time()
elapsed_time = end_time - start_time
log_to_db(
handler.client_ip,
handler.token,
handler.prompt,
generated_text,
elapsed_time,
handler.parameters,
r_headers,
response_status_code,
r_url,
handler.backend_url,
)
log_to_db(
handler.client_ip,
handler.token,
handler.prompt,
generated_text,
elapsed_time,
handler.parameters,
r_headers,
response_status_code,
r_url,
handler.backend_url,
)
finally:
if event_id:
redis.publish(event_id, 'finished')
else:
print('event_id was None!')
return Response(generate(), mimetype='text/event-stream')
except Exception:
traceback.print_exc()
return 'INTERNAL SERVER', 500
finally:
if event_id:
redis.publish(event_id, 'finished')
else:
print('event_id was None!')

64
llm_server/routes/queue.py
View File

@ -1,10 +1,12 @@
import json
import pickle
import time
from typing import Tuple
from uuid import uuid4
from redis import Redis
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
@ -20,7 +22,7 @@ def decrement_ip_count(client_ip: str, redis_key):
class RedisPriorityQueue:
def __init__(self, name: str = 'priority_queue', db: int = 12):
def __init__(self, name, db: int = 12):
self.redis = RedisCustom(name, db=db)
def put(self, item, priority, selected_model):
@ -98,9 +100,6 @@ class DataEvent:
return pickle.loads(item['data'])
priority_queue = RedisPriorityQueue()
def update_active_workers(key: str, operation: str):
if operation == 'incr':
redis.incr(f'active_gen_workers:{key}')
@ -118,3 +117,60 @@ def incr_active_workers(selected_model: str, backend_url: str):
def decr_active_workers(selected_model: str, backend_url: str):
update_active_workers(selected_model, 'decr')
update_active_workers(backend_url, 'decr')
class PriorityQueue:
def __init__(self, backends: list = 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.lpush('backends', item)
def get_backends(self):
return [x.decode('utf-8') for x in self.redis.lrange('backends', 0, -1)]
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_queued_ip_count(client_ip)
return count
def put(self, backend_url, item: Tuple[dict, str, str, dict], priority: int, selected_model: str):
queue = RedisPriorityQueue(backend_url)
return queue.put(item, priority, selected_model)
def len(self, model_name):
count = 0
backends_with_models = []
for k in self.get_backends():
info = cluster_config.get_backend(k)
if info.get('model') == model_name:
backends_with_models.append(k)
for backend_url in backends_with_models:
queue = RedisPriorityQueue(backend_url)
count += queue.len(model_name)
return count
def __len__(self):
count = 0
for backend_url in self.get_backends():
queue = RedisPriorityQueue(backend_url)
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()

4
llm_server/routes/request_handler.py
View File

@ -7,7 +7,7 @@ from flask import Response, request
from llm_server import opts
from llm_server.cluster.cluster_config import cluster_config, get_a_cluster_backend
from llm_server.custom_redis import redis
from llm_server.database.database import get_token_ratelimit, do_db_log
from llm_server.database.database import get_token_ratelimit
from llm_server.database.log_to_db import log_to_db
from llm_server.helpers import auto_set_base_client_api
from llm_server.llm.oobabooga.ooba_backend import OobaboogaBackend
@ -131,7 +131,7 @@ class RequestHandler:
request_valid, invalid_response = self.validate_request(prompt, do_log=True)
if not request_valid:
return (False, None, None, 0), invalid_response
event = priority_queue.put((llm_request, self.client_ip, self.token, self.parameters, self.backend_url), self.token_priority, self.selected_model)
event = priority_queue.put(self.backend_url, (llm_request, self.client_ip, self.token, self.parameters), self.token_priority, self.selected_model)
else:
event = None

2
llm_server/routes/v1/generate_stream.py
View File

@ -122,7 +122,7 @@ def do_stream(ws, model_name):
event = None
if not handler.is_client_ratelimited():
# Add a dummy event to the queue and wait for it to reach a worker
event = priority_queue.put((None, handler.client_ip, handler.token, None, handler.backend_url), handler.token_priority, handler.selected_model)
event = priority_queue.put(handler.backend_url, (None, handler.client_ip, handler.token, None), handler.token_priority, handler.selected_model)
if not event:
log_to_db(
handler.client_ip,

45
llm_server/workers/inferencer.py
View File

@ -1,29 +1,24 @@
import threading
import time
from uuid import uuid4
from llm_server.cluster.cluster_config import cluster_config, get_a_cluster_backend
from llm_server.custom_redis import redis
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, priority_queue
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():
def worker(backend_url):
queue = RedisPriorityQueue(backend_url)
while True:
(request_json_body, client_ip, token, parameters, backend_url), event_id, selected_model = priority_queue.get()
if not backend_url:
backend_url = get_a_cluster_backend(selected_model)
else:
backend_url = cluster_config.validate_backend(backend_url)
(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)
need_to_wait(backend_url)
try:
if not request_json_body:
# This was a dummy request from the streaming handlers.
@ -37,7 +32,6 @@ def worker():
redis.publish(event_id, 'begin')
for item in pubsub.listen():
if item['type'] == 'message' and item['data'].decode('utf-8') == 'finished':
# Once the handler is complete, move on.
break
time.sleep(0.1)
else:
@ -50,23 +44,12 @@ def worker():
decr_active_workers(selected_model, backend_url)
def start_workers(num_workers: int):
def start_workers(cluster: dict):
i = 0
for _ in range(num_workers):
t = threading.Thread(target=worker)
t.daemon = True
t.start()
i += 1
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.')
def need_to_wait(backend_url: str):
# We need to check the number of active workers since the streaming endpoint may be doing something.
active_workers = redis.get(f'active_gen_workers:{backend_url}', 0, dtype=int)
concurrent_gens = cluster_config.get_backend(backend_url).get('concurrent_gens', 1)
s = time.time()
while active_workers >= concurrent_gens:
time.sleep(0.01)
e = time.time()
if e - s > 0.1:
print(f'Worker was delayed {e - s} seconds.')

2
llm_server/workers/printer.py
View File

@ -25,6 +25,4 @@ def console_printer():
processing_count += redis.get(k, default=0, dtype=int)
backends = [k for k, v in cluster_config.all().items() if v['online']]
logger.info(f'REQUEST QUEUE -> Processing: {processing_count} | Queued: {len(priority_queue)} | Backends Online: {len(backends)}')
priority_queue.print_all_items()
print('============================')
time.sleep(1)

2
llm_server/workers/threader.py
View File

@ -20,7 +20,7 @@ def cache_stats():
def start_background():
start_workers(opts.cluster_workers)
start_workers(opts.cluster)
t = Thread(target=main_background_thread)
t.daemon = True