diff --git a/llm_server/custom_redis.py b/llm_server/custom_redis.py index 886230a..a055537 100644 --- a/llm_server/custom_redis.py +++ b/llm_server/custom_redis.py @@ -16,7 +16,8 @@ ONE_MONTH_SECONDS = 2678000 class RedisCustom(Redis): """ - A wrapper class to set prefixes to keys. + A simple wrapper class for Redis to create a "namespace" within a DB, + which simplyifies key management. """ def __init__(self, prefix, **kwargs): @@ -233,21 +234,22 @@ class RedisCustom(Redis): return self.redis.lpop(self._key(name), count) def zrange( - self, - name: KeyT, - start: int, - end: int, - desc: bool = False, - withscores: bool = False, - score_cast_func: Union[type, Callable] = float, - byscore: bool = False, - bylex: bool = False, - offset: int = None, - num: int = None, + self, + name: KeyT, + start: int, + end: int, + desc: bool = False, + withscores: bool = False, + score_cast_func: Union[type, Callable] = float, + byscore: bool = False, + bylex: bool = False, + offset: int = None, + num: int = None, ): return self.redis.zrange(self._key(name), start, end, desc, withscores, score_cast_func, byscore, bylex, offset, num) def zrem(self, name: KeyT, *values: FieldT): return self.redis.zrem(self._key(name), *values) + redis = RedisCustom('local_llm') diff --git a/llm_server/routes/queue.py b/llm_server/routes/queue.py index 62c9c42..3bbdd1c 100644 --- a/llm_server/routes/queue.py +++ b/llm_server/routes/queue.py @@ -23,6 +23,10 @@ def decrement_ip_count(client_ip: str, redis_key): class RedisPriorityQueue: + """ + A queue for a specific backend. + """ + def __init__(self, name, db: int = 12): self.name = name self.redis = RedisCustom(name, db=db) @@ -99,6 +103,10 @@ class RedisPriorityQueue: class DataEvent: + """ + Class to simplify pub/sub communication between consumers and producers (MASTERS and SLAVES lololololol). + """ + def __init__(self, event_id=None): self.event_id = event_id if event_id else str(uuid4()) self.redis = Redis(host='localhost', port=6379, db=14) @@ -134,6 +142,10 @@ def decr_active_workers(selected_model: str, backend_url: str): class PriorityQueue: + """ + Helper class to wrangler all the different queues. + """ + def __init__(self, backends: set = None): """ Only have to load the backends once. diff --git a/llm_server/routes/server_error.py b/llm_server/routes/server_error.py index fec3836..b18d7f1 100644 --- a/llm_server/routes/server_error.py +++ b/llm_server/routes/server_error.py @@ -1,3 +1,3 @@ def handle_server_error(e): print(e) - return {'error': True}, 500 + return {'error': True, 'code': 500, 'message': 'Internal Server Error :('}, 500