fix import error
This commit is contained in:
parent
e07e31df0a
commit
9b819573e8
|
@ -1,9 +1,68 @@
|
|||
import hashlib
|
||||
import pickle
|
||||
|
||||
from llm_server import opts
|
||||
from llm_server.cluster.redis_config_cache import RedisClusterStore
|
||||
from llm_server.cluster.redis_cycle import add_backend_cycler, redis_cycle
|
||||
from llm_server.cluster.stores import redis_running_models
|
||||
from llm_server.custom_redis import RedisCustom
|
||||
from llm_server.routes.helpers.model import estimate_model_size
|
||||
|
||||
|
||||
class RedisClusterStore:
|
||||
def __init__(self, name: str, **kwargs):
|
||||
self.name = name
|
||||
self.config_redis = RedisCustom(name, **kwargs)
|
||||
|
||||
def clear(self):
|
||||
self.config_redis.flush()
|
||||
|
||||
def load(self, config: dict):
|
||||
for k, v in config.items():
|
||||
self.add_backend(k, v)
|
||||
|
||||
def add_backend(self, name: str, values: dict):
|
||||
self.config_redis.hset(name, mapping={k: pickle.dumps(v) for k, v in values.items()})
|
||||
self.set_backend_value(name, 'online', False)
|
||||
h = hashlib.sha256(name.encode('utf-8')).hexdigest()
|
||||
self.set_backend_value(name, 'hash', f'{h[:8]}-{h[-8:]}')
|
||||
|
||||
def set_backend_value(self, backend: str, key: str, value):
|
||||
# By storing the value as a pickle we don't have to cast anything when getting the value from Redis.
|
||||
self.config_redis.hset(backend, key, pickle.dumps(value))
|
||||
|
||||
def get_backend(self, name: str):
|
||||
r = self.config_redis.hgetall(name)
|
||||
output = {}
|
||||
for k, v in r.items():
|
||||
output[k.decode('utf8')] = pickle.loads(v)
|
||||
return output
|
||||
|
||||
def all(self):
|
||||
keys = self.config_redis.keys('*')
|
||||
if keys:
|
||||
result = {}
|
||||
for key in keys:
|
||||
if key != f'{self.name}:____':
|
||||
v = self.get_backend(key)
|
||||
result[key] = v
|
||||
return result
|
||||
else:
|
||||
return {}
|
||||
|
||||
def validate_backend(self, backend_url: str):
|
||||
"""
|
||||
Returns the backend URL that was given, or a new one if that was offline.
|
||||
:param backend_url:
|
||||
:return:
|
||||
"""
|
||||
backend_info = self.get_backend(backend_url)
|
||||
if not backend_info['online']:
|
||||
old = backend_url
|
||||
backend_url = get_a_cluster_backend()
|
||||
print(f'Backend {old} offline. Request was redirected to {backend_url}')
|
||||
return backend_url
|
||||
|
||||
|
||||
cluster_config = RedisClusterStore('cluster_config')
|
||||
|
||||
|
||||
|
|
|
@ -1,60 +0,0 @@
|
|||
import hashlib
|
||||
import pickle
|
||||
|
||||
from llm_server.cluster.cluster_config import get_a_cluster_backend
|
||||
from llm_server.custom_redis import RedisCustom
|
||||
|
||||
|
||||
class RedisClusterStore:
|
||||
def __init__(self, name: str, **kwargs):
|
||||
self.name = name
|
||||
self.config_redis = RedisCustom(name, **kwargs)
|
||||
|
||||
def clear(self):
|
||||
self.config_redis.flush()
|
||||
|
||||
def load(self, config: dict):
|
||||
for k, v in config.items():
|
||||
self.add_backend(k, v)
|
||||
|
||||
def add_backend(self, name: str, values: dict):
|
||||
self.config_redis.hset(name, mapping={k: pickle.dumps(v) for k, v in values.items()})
|
||||
self.set_backend_value(name, 'online', False)
|
||||
h = hashlib.sha256(name.encode('utf-8')).hexdigest()
|
||||
self.set_backend_value(name, 'hash', f'{h[:8]}-{h[-8:]}')
|
||||
|
||||
def set_backend_value(self, backend: str, key: str, value):
|
||||
# By storing the value as a pickle we don't have to cast anything when getting the value from Redis.
|
||||
self.config_redis.hset(backend, key, pickle.dumps(value))
|
||||
|
||||
def get_backend(self, name: str):
|
||||
r = self.config_redis.hgetall(name)
|
||||
output = {}
|
||||
for k, v in r.items():
|
||||
output[k.decode('utf8')] = pickle.loads(v)
|
||||
return output
|
||||
|
||||
def all(self):
|
||||
keys = self.config_redis.keys('*')
|
||||
if keys:
|
||||
result = {}
|
||||
for key in keys:
|
||||
if key != f'{self.name}:____':
|
||||
v = self.get_backend(key)
|
||||
result[key] = v
|
||||
return result
|
||||
else:
|
||||
return {}
|
||||
|
||||
def validate_backend(self, backend_url: str):
|
||||
"""
|
||||
Returns the backend URL that was given, or a new one if that was offline.
|
||||
:param backend_url:
|
||||
:return:
|
||||
"""
|
||||
backend_info = self.get_backend(backend_url)
|
||||
if not backend_info['online']:
|
||||
old = backend_url
|
||||
backend_url = get_a_cluster_backend()
|
||||
print(f'Backend {old} offline. Request was redirected to {backend_url}')
|
||||
return backend_url
|
Reference in New Issue