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')

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.

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