118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
import numpy as np
|
|
|
|
from llm_server import opts
|
|
from llm_server.cluster.cluster_config import cluster_config, get_a_cluster_backend
|
|
from llm_server.cluster.stores import redis_running_models
|
|
from llm_server.custom_redis import redis
|
|
from llm_server.llm.generator import generator
|
|
from llm_server.llm.info import get_info
|
|
from llm_server.llm.vllm.vllm_backend import VLLMBackend
|
|
from llm_server.routes.queue import priority_queue
|
|
from llm_server.routes.stats import calculate_wait_time, get_active_gen_workers_model
|
|
|
|
|
|
def get_backends_from_model(model_name: str):
|
|
return [x.decode('utf-8') for x in redis_running_models.smembers(model_name)]
|
|
|
|
|
|
def get_running_models():
|
|
return redis_running_models.keys()
|
|
|
|
|
|
def purge_backend_from_running_models(backend_url: str):
|
|
keys = redis_running_models.keys()
|
|
pipeline = redis_running_models.pipeline()
|
|
for model in keys:
|
|
pipeline.srem(model, backend_url)
|
|
pipeline.execute()
|
|
|
|
|
|
def is_valid_model(model_name: str):
|
|
return redis_running_models.exists(model_name)
|
|
|
|
|
|
def test_backend(backend_url: str, test_prompt: bool = False):
|
|
backend_info = cluster_config.get_backend(backend_url)
|
|
if test_prompt:
|
|
handler = VLLMBackend(backend_url)
|
|
parameters, _ = handler.get_parameters({
|
|
"stream": False,
|
|
"temperature": 0,
|
|
"max_new_tokens": 3,
|
|
})
|
|
data = {
|
|
'prompt': 'test prompt',
|
|
**parameters
|
|
}
|
|
try:
|
|
success, response, err = generator(data, backend_url, timeout=10)
|
|
if not success or not response or err:
|
|
return False, {}
|
|
except:
|
|
return False, {}
|
|
i = get_info(backend_url, backend_info['mode'])
|
|
if not i.get('model'):
|
|
return False, {}
|
|
return True, i
|
|
|
|
|
|
def get_model_choices(regen: bool = False):
|
|
if not regen:
|
|
c = redis.getp('model_choices')
|
|
if c:
|
|
return c
|
|
|
|
base_client_api = redis.get('base_client_api', dtype=str)
|
|
running_models = get_running_models()
|
|
model_choices = {}
|
|
for model in running_models:
|
|
b = get_backends_from_model(model)
|
|
|
|
context_size = []
|
|
avg_gen_per_worker = []
|
|
concurrent_gens = 0
|
|
for backend_url in b:
|
|
backend_info = cluster_config.get_backend(backend_url)
|
|
if backend_info.get('model_config'):
|
|
context_size.append(backend_info['model_config']['max_position_embeddings'])
|
|
if backend_info.get('average_generation_elapsed_sec'):
|
|
avg_gen_per_worker.append(backend_info['average_generation_elapsed_sec'])
|
|
concurrent_gens += backend_info['concurrent_gens']
|
|
|
|
active_gen_workers = get_active_gen_workers_model(model)
|
|
proompters_in_queue = priority_queue.len(model)
|
|
|
|
if len(avg_gen_per_worker):
|
|
average_generation_elapsed_sec = np.average(avg_gen_per_worker)
|
|
else:
|
|
average_generation_elapsed_sec = 0
|
|
estimated_wait_sec = calculate_wait_time(average_generation_elapsed_sec, proompters_in_queue, concurrent_gens, active_gen_workers)
|
|
|
|
model_choices[model] = {
|
|
'model': model,
|
|
'client_api': f'https://{base_client_api}/{model}',
|
|
'ws_client_api': f'wss://{base_client_api}/{model}/v1/stream' if opts.enable_streaming else None,
|
|
'openai_client_api': f'https://{base_client_api}/openai/{model}/v1' if opts.enable_openi_compatible_backend else 'disabled',
|
|
'backend_count': len(b),
|
|
'estimated_wait': estimated_wait_sec,
|
|
'queued': proompters_in_queue,
|
|
'processing': active_gen_workers,
|
|
'avg_generation_time': average_generation_elapsed_sec,
|
|
'concurrent_gens': concurrent_gens
|
|
}
|
|
|
|
if len(context_size):
|
|
model_choices[model]['context_size'] = min(context_size)
|
|
|
|
# Python wants to sort lowercase vs. uppercase letters differently.
|
|
model_choices = dict(sorted(model_choices.items(), key=lambda item: item[0].upper()))
|
|
|
|
default_backend_url = get_a_cluster_backend()
|
|
default_backend_info = cluster_config.get_backend(default_backend_url)
|
|
if not default_backend_info.get('model'):
|
|
return {}, None
|
|
default_model = default_backend_info['model']
|
|
|
|
redis.setp('model_choices', (model_choices, default_model))
|
|
return model_choices, default_model
|