50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
import hashlib
|
|
import pickle
|
|
|
|
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 get(self, name: str):
|
|
# return self.all().get(name)
|