Computing message queue positions (LAION-AI#2235)

Introduces counters for work queues that allows us to track the
positions of enqueued work requests without having to iterate through
the queues

Author: yk

inference/server/oasst_inference_server/chat_repository.py

@@ -49,7 +49,6 @@ async def start_work(
         message.state = inference.MessageState.in_progress
         message.work_begin_at = datetime.datetime.utcnow()
         message.worker_id = worker_id
-        message.worker_compat_hash = worker_config.compat_hash
         message.worker_config = worker_config
         await self.session.commit()
         logger.debug(f"Started work on message {message_id}")

inference/server/oasst_inference_server/queueing.py

@@ -3,36 +3,76 @@
 
 
 class RedisQueue:
-    def __init__(self, redis_client: redis.Redis, queue_id: str) -> None:
+    def __init__(
+        self,
+        redis_client: redis.Redis,
+        queue_id: str,
+        expire: int | None = None,
+        with_counter: bool = False,
+        counter_pos_expire: int = 1,
+    ) -> None:
         self.redis_client = redis_client
         self.queue_id = queue_id
+        self.expire = expire
+        self.with_counter = with_counter
+        self.counter_pos_expire = counter_pos_expire
 
-    async def enqueue(self, value: str, expire: int | None = None) -> None:
-        pushed = await self.redis_client.rpush(self.queue_id, value)
-        if expire is not None:
-            await self.set_expire(expire)
-        return pushed
+    async def enqueue(self, value: str) -> int | None:
+        await self.redis_client.rpush(self.queue_id, value)
+        if self.expire is not None:
+            await self.set_expire(self.expire)
+        if self.with_counter:
+            ctr = await self.redis_client.incr(f"ctr_enq:{self.queue_id}")
+            await self.redis_client.set(f"pos:{value}", ctr, ex=self.counter_pos_expire)
+        else:
+            ctr = None
+        return ctr
 
-    async def dequeue(self, timeout: int = 1) -> str:
-        return await self.redis_client.blpop(self.queue_id, timeout=timeout)
+    async def dequeue(self, timeout: int = 1) -> str | None:
+        val = await self.redis_client.blpop(self.queue_id, timeout=timeout)
+        if val is not None and self.with_counter:
+            await self.redis_client.incr(f"ctr_deq:{self.queue_id}")
+        return val
 
     async def set_expire(self, timeout: int) -> None:
         return await self.redis_client.expire(self.queue_id, timeout)
 
+    async def get_enq_counter(self) -> int:
+        if not self.with_counter:
+            return 0
+        enq = await self.redis_client.get(f"ctr_enq:{self.queue_id}")
+        enq = int(enq) if enq is not None else 0
+        return enq
 
-def chat_queue(redis_client: redis.Redis, chat_id: str) -> RedisQueue:
-    return RedisQueue(redis_client, f"chat:{chat_id}")
+    async def get_deq_counter(self) -> int:
+        if not self.with_counter:
+            return 0
+        deq = await self.redis_client.get(f"ctr_deq:{self.queue_id}")
+        deq = int(deq) if deq is not None else 0
+        return deq
+
+    async def get_length(self) -> int:
+        return await self.redis_client.llen(self.queue_id)
+
+
+async def get_pos_value(redis_client: redis.Redis, message_id: str) -> int:
+    val = await redis_client.get(f"pos:{message_id}")
+    if val is None:
+        return 0
+    return int(val)
 
 
 def message_queue(redis_client: redis.Redis, message_id: str) -> RedisQueue:
-    return RedisQueue(redis_client, f"message:{message_id}")
+    return RedisQueue(redis_client, f"message:{message_id}", expire=settings.message_queue_expire)
 
 
 def work_queue(redis_client: redis.Redis, worker_compat_hash: str) -> RedisQueue:
     if settings.allowed_worker_compat_hashes != "*":
         if worker_compat_hash not in settings.allowed_worker_compat_hashes_list:
             raise ValueError(f"Worker compat hash {worker_compat_hash} not allowed")
-    return RedisQueue(redis_client, f"work:{worker_compat_hash}")
+    return RedisQueue(
+        redis_client, f"work:{worker_compat_hash}", with_counter=True, counter_pos_expire=settings.message_queue_expire
+    )
 
 
 def compliance_queue(redis_client: redis.Redis, worker_id: str) -> RedisQueue:

inference/server/oasst_inference_server/routes/chats.py

@@ -1,3 +1,5 @@
+import asyncio
+
 import fastapi
 import pydantic
 from fastapi import Depends
@@ -94,6 +96,7 @@ async def create_assistant_message(
             assistant_message = await ucr.initiate_assistant_message(
                 parent_id=request.parent_id,
                 work_parameters=work_parameters,
+                worker_compat_hash=model_config.compat_hash,
             )
         queue = queueing.work_queue(deps.redis_client, model_config.compat_hash)
         logger.debug(f"Adding {assistant_message.id=} to {queue.queue_id} for {chat_id}")
@@ -133,19 +136,35 @@ async def message_events(
     if message.has_finished:
         raise fastapi.HTTPException(status_code=204, detail=message.state)
 
-    async def event_generator(chat_id: str, message_id: str):
+    async def event_generator(chat_id: str, message_id: str, worker_compat_hash: str | None):
         redis_client = deps.make_redis_client()
-        queue = queueing.message_queue(redis_client, message_id=message_id)
+        message_queue = queueing.message_queue(redis_client, message_id=message_id)
+        work_queue = (
+            queueing.work_queue(redis_client, worker_compat_hash=worker_compat_hash)
+            if worker_compat_hash is not None
+            else None
+        )
         has_started = False
         try:
             while True:
-                item = await queue.dequeue(timeout=settings.pending_event_interval)
+                item = await message_queue.dequeue(timeout=settings.pending_event_interval)
                 if item is None:
                     if not has_started:
+                        if work_queue is None:
+                            qpos, qlen = 0, 1
+                        else:
+                            # TODO: make more efficient, e.g. pipeline
+                            [qdeq, qenq, mpos] = await asyncio.gather(
+                                work_queue.get_deq_counter(),
+                                work_queue.get_enq_counter(),
+                                queueing.get_pos_value(redis_client, message_id),
+                            )
+                            qpos = max(mpos - qdeq, 0)
+                            qlen = max(qenq - qdeq, qpos)
                         yield {
                             "data": chat_schema.PendingResponseEvent(
-                                queue_position=0,
-                                queue_size=1,
+                                queue_position=qpos,
+                                queue_size=qlen,
                             ).json()
                         }
                     continue
@@ -188,7 +207,9 @@ async def event_generator(chat_id: str, message_id: str):
         finally:
             await redis_client.close()
 
-    return EventSourceResponse(event_generator(chat_id=chat_id, message_id=message_id))
+    return EventSourceResponse(
+        event_generator(chat_id=chat_id, message_id=message_id, worker_compat_hash=message.worker_compat_hash)
+    )
 
 
 @router.post("/{chat_id}/messages/{message_id}/votes")

inference/server/oasst_inference_server/routes/workers.py

@@ -327,7 +327,7 @@ async def handle_token_response(
         deps.redis_client,
         message_id=work_response_container.message_id,
     )
-    await message_queue.enqueue(response.json(), expire=settings.message_queue_expire)
+    await message_queue.enqueue(response.json())
     work_response_container.num_responses += 1
 
 
@@ -352,7 +352,7 @@ async def handle_generated_text_response(
             deps.redis_client,
             message_id=message_id,
         )
-        await message_queue.enqueue(message_packet.json(), expire=settings.message_queue_expire)
+        await message_queue.enqueue(message_packet.json())
     finally:
         del work_request_map[response.request_id]
 
@@ -365,7 +365,7 @@ async def abort_message(message_id: str, error: str):
         deps.redis_client,
         message_id=message_id,
     )
-    await message_queue.enqueue(response.json(), expire=settings.message_queue_expire)
+    await message_queue.enqueue(response.json())
 
 
 async def handle_error_response(
@@ -396,4 +396,4 @@ async def handle_timeout(message: inference.MessageRead):
         deps.redis_client,
         message_id=message.id,
     )
-    await message_queue.enqueue(response.json(), expire=settings.message_queue_expire)
+    await message_queue.enqueue(response.json())

inference/server/oasst_inference_server/user_chat_repository.py

@@ -110,7 +110,7 @@ async def add_prompter_message(self, chat_id: str, parent_id: str | None, conten
         return message
 
     async def initiate_assistant_message(
-        self, parent_id: str, work_parameters: inference.WorkParameters
+        self, parent_id: str, work_parameters: inference.WorkParameters, worker_compat_hash: str
     ) -> models.DbMessage:
         logger.info(f"Adding stub assistant message to {parent_id=}")
 
@@ -154,6 +154,7 @@ async def initiate_assistant_message(
             parent_id=parent_id,
             state=inference.MessageState.pending,
             work_parameters=work_parameters,
+            worker_compat_hash=worker_compat_hash,
         )
         self.session.add(message)
         await self.session.commit()

inference/text-client/__main__.py

@@ -10,13 +10,13 @@
 
 
 @app.command()
-def main(backend_url: str = "http://127.0.0.1:8000", model_config_name="distilgpt2"):
+def main(backend_url: str = "http://127.0.0.1:8000", model_config_name="distilgpt2", username="test1"):
     """Simple REPL client."""
     while True:
         try:
             # login
             client = utils.DebugClient(backend_url)
-            client.login("test1")
+            client.login(username)
             chat_id = client.create_chat()
             typer.echo(f"Chat ID: {chat_id}")
             while True: