109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
from llm_server import opts
|
|
from llm_server.cluster.cluster_config import cluster_config
|
|
from llm_server.cluster.redis_cycle import add_backend_cycler, redis_cycle
|
|
from llm_server.cluster.stores import redis_running_models
|
|
from llm_server.llm.generator import generator
|
|
from llm_server.llm.info import get_info
|
|
from llm_server.routes.helpers.model import estimate_model_size
|
|
|
|
|
|
def test_backend(backend_url: str, test_prompt: bool = False):
|
|
backend_info = cluster_config.get_backend(backend_url)
|
|
if test_prompt:
|
|
data = {
|
|
"prompt": "Test prompt",
|
|
"stream": False,
|
|
"temperature": 0,
|
|
"max_new_tokens": 3,
|
|
}
|
|
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_backends():
|
|
backends = cluster_config.all()
|
|
result = {}
|
|
for k, v in backends.items():
|
|
b = cluster_config.get_backend(k)
|
|
status = b.get('online', False)
|
|
priority = b['priority']
|
|
result[k] = {'status': status, 'priority': priority}
|
|
|
|
if not opts.prioritize_by_size:
|
|
online_backends = sorted(
|
|
((url, info) for url, info in backends.items() if info['online']),
|
|
key=lambda kv: -kv[1]['priority'],
|
|
reverse=True
|
|
)
|
|
else:
|
|
online_backends = sorted(
|
|
((url, info) for url, info in backends.items() if info['online']),
|
|
key=lambda kv: estimate_model_size(kv[1]['model_config']),
|
|
reverse=True
|
|
)
|
|
offline_backends = sorted(
|
|
((url, info) for url, info in backends.items() if not info['online']),
|
|
key=lambda kv: -kv[1]['priority'],
|
|
reverse=True
|
|
)
|
|
return [url for url, info in online_backends], [url for url, info in offline_backends]
|
|
|
|
|
|
def get_a_cluster_backend(model=None):
|
|
"""
|
|
Get a backend from Redis. If there are no online backends, return None.
|
|
If `model` is not supplied, we will pick one ourself.
|
|
"""
|
|
if model:
|
|
# First, determine if there are multiple backends hosting the same model.
|
|
backends_hosting_model = [i.decode('utf-8') for i in redis_running_models.smembers(model)]
|
|
|
|
# If so, create an iterator for those backends
|
|
if len(backends_hosting_model):
|
|
add_backend_cycler(model, backends_hosting_model)
|
|
cycled = redis_cycle(model)
|
|
if len(cycled):
|
|
return cycled[0]
|
|
else:
|
|
# No backend hosting that model
|
|
return None
|
|
else:
|
|
online, _ = get_backends()
|
|
if len(online):
|
|
return online[0]
|
|
|
|
|
|
def get_backends_from_model(model_name: str):
|
|
return [x.decode('utf-8') for x in redis_running_models.smembers(model_name)]
|
|
|
|
|
|
# def verify_context_size(model_name:str):
|
|
# b = get_backends_from_model(model_name)
|
|
# for backend_url in b:
|
|
# backend_info = cluster_config.get_backend(backend_url)
|
|
# backend_info.get()
|
|
|
|
|
|
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)
|