unified queueing

inference/server/main.py

@@ -12,7 +12,7 @@ import websockets.exceptions
 from fastapi import Depends
 from fastapi.middleware.cors import CORSMiddleware
 from loguru import logger
-from oasst_inference_server import interface
+from oasst_inference_server import interface, queueing
 from oasst_inference_server.chat_repository import ChatRepository
 from oasst_inference_server.database import db_engine
 from oasst_inference_server.settings import settings
@@ -40,7 +40,7 @@ app.add_middleware(
 
 
 # create async redis client
-redisClient = redis.Redis(
+redis_client = redis.Redis(
     host=settings.redis_host, port=settings.redis_port, db=settings.redis_db, decode_responses=True
 )
 
@@ -55,6 +55,12 @@ def create_chat_repository(session: sqlmodel.Session = Depends(create_session)):
     return repository
 
 
+@contextlib.contextmanager
+def manual_chat_repository():
+    with contextlib.contextmanager(create_session)() as session:
+        yield create_chat_repository(session)
+
+
 if settings.update_alembic:
 
     @app.on_event("startup")
@@ -105,9 +111,9 @@ async def get_chat(id: str, chat_repository: ChatRepository = Depends(create_cha
     return chat
 
 
-@app.post("/chat/{id}/message")
+@app.post("/chat/{chat_id}/message")
 async def create_message(
-    id: str,
+    chat_id: str,
     message_request: interface.MessageRequest,
     fastapi_request: fastapi.Request,
     chat_repository: ChatRepository = Depends(create_chat_repository),
@@ -115,21 +121,24 @@ async def create_message(
     """Allows the client to stream the results of a request."""
 
     try:
-        chat_repository.add_prompter_message(id=id, message_request=message_request)
+        chat_repository.add_prompter_message(chat_id=chat_id, message_request=message_request)
+        queue = queueing.work_queue(redis_client, message_request.worker_compat_hash)
+        logger.debug(f"Adding {chat_id} to {queue.queue_id}")
+        await queue.enqueue(chat_id)
+        logger.debug(f"Added message to {queue.queue_id} for {chat_id}")
     except Exception:
         logger.exception("Error adding prompter message")
         return fastapi.Response(status_code=500)
 
-    async def event_generator(id):
+    async def event_generator(chat_id):
+        queue = queueing.chat_queue(redis_client, chat_id)
         result_data = []
-
         try:
             while True:
                 if await fastapi_request.is_disconnected():
                     logger.warning("Client disconnected")
                     return
-
-                item = await redisClient.blpop(id, 1)
+                item = await queue.dequeue()
                 if item is None:
                     continue
 
@@ -144,25 +153,26 @@ async def create_message(
                     "retry": settings.sse_retry_timeout,
                     "data": interface.TokenResponseEvent(token=response_packet.token).json(),
                 }
-            logger.info(f"Finished streaming {id} {len(result_data)=}")
+            logger.info(f"Finished streaming {chat_id} {len(result_data)=}")
         except Exception:
-            logger.exception(f"Error streaming {id}")
+            logger.exception(f"Error streaming {chat_id}")
             raise
 
         try:
-            with contextlib.contextmanager(create_session)() as session:
-                chat_repository = create_chat_repository(session)
-                chat_repository.add_assistant_message(id=id, text=response_packet.generated_text.text)
+            with manual_chat_repository() as chat_repository:
+                chat_repository.add_assistant_message(chat_id=chat_id, text=response_packet.generated_text.text)
         except Exception:
             logger.exception("Error adding assistant message")
 
-    return EventSourceResponse(event_generator(id))
+    return EventSourceResponse(event_generator(chat_id))
 
 
 @app.websocket("/work")
-async def work(websocket: fastapi.WebSocket, chat_repository: ChatRepository = Depends(create_chat_repository)):
+async def work(websocket: fastapi.WebSocket):
     await websocket.accept()
     worker_config = inference.WorkerConfig.parse_raw(await websocket.receive_text())
+    queue_id = f"work:{worker_config.compat_hash}"
+    work_queue = queueing.RedisQueue(redis_client, queue_id)
     try:
         while True:
             if websocket.client_state == fastapi.websockets.WebSocketState.DISCONNECTED:
@@ -171,55 +181,62 @@ async def work(websocket: fastapi.WebSocket, chat_repository: ChatRepository = D
             # find a pending task that matches the worker's config
             # could also be implemented using task queues
             # but general compatibility matching is tricky
-            for chat in chat_repository.get_pending_chats():
-                request = chat.pending_message_request
-                if request.compatible_with(worker_config):
-                    break
-            else:
+            item = await work_queue.dequeue()
+            if item is None:
                 await asyncio.sleep(1)
                 continue
+            else:
+                _, chat_id = item
 
-            chat_repository.set_chat_state(chat.id, interface.MessageRequestState.in_progress)
+            with manual_chat_repository() as chat_repository:
+                chat = chat_repository.get_chat_by_id(chat_id)
+                request = chat.pending_message_request
 
-            work_request = inference.WorkRequest(
-                conversation=chat.conversation,
-                model_name=request.model_name,
-                max_new_tokens=request.max_new_tokens,
-            )
+                chat_repository.set_chat_state(chat.id, interface.MessageRequestState.in_progress)
 
-            logger.info(f"Created {work_request=}")
-            try:
-                await websocket.send_text(work_request.json())
-            except websockets.exceptions.ConnectionClosedError:
-                logger.warning("Worker disconnected")
-                websocket.close()
-                chat_repository.set_chat_state(chat.id, interface.MessageRequestState.pending)
-                break
+                work_request = inference.WorkRequest(
+                    conversation=chat.conversation,
+                    model_name=request.model_name,
+                    max_new_tokens=request.max_new_tokens,
+                )
 
-            logger.debug(f"Sent {work_request=} to worker.")
-
-            try:
-                in_progress = False
-                while True:
-                    # maybe unnecessary to parse and re-serialize
-                    # could just pass the raw string and mark end via empty string
-                    response_packet = inference.WorkResponsePacket.parse_raw(await websocket.receive_text())
-                    in_progress = True
-                    await redisClient.rpush(chat.id, response_packet.json())
-                    if response_packet.is_end:
-                        logger.debug(f"Received {response_packet=} from worker. Ending.")
-                        break
-            except fastapi.WebSocketException:
-                # TODO: handle this better
-                logger.exception(f"Websocket closed during handling of {chat.id}")
-                if in_progress:
-                    logger.warning(f"Aborting {chat.id=}")
-                    chat_repository.set_chat_state(chat.id, interface.MessageRequestState.aborted_by_worker)
-                else:
-                    logger.warning(f"Marking {chat.id=} as pending since no work was done.")
+                logger.info(f"Created {work_request=}")
+                try:
+                    await websocket.send_text(work_request.json())
+                except websockets.exceptions.ConnectionClosedError:
+                    logger.warning("Worker disconnected")
+                    websocket.close()
                     chat_repository.set_chat_state(chat.id, interface.MessageRequestState.pending)
-                raise
+                    await work_queue.enqueue(chat.id)
+                    break
 
-            chat_repository.set_chat_state(chat.id, interface.MessageRequestState.complete)
+                logger.debug(f"Sent {work_request=} to worker.")
+
+                chat_queue = queueing.chat_queue(redis_client, chat.id)
+
+                try:
+                    in_progress = False
+                    while True:
+                        # maybe unnecessary to parse and re-serialize
+                        # could just pass the raw string and mark end via empty string
+                        response_packet = inference.WorkResponsePacket.parse_raw(await websocket.receive_text())
+                        in_progress = True
+                        await chat_queue.enqueue(response_packet.json())
+                        if response_packet.is_end:
+                            logger.debug(f"Received {response_packet=} from worker. Ending.")
+                            break
+                except fastapi.WebSocketException:
+                    # TODO: handle this better
+                    logger.exception(f"Websocket closed during handling of {chat.id}")
+                    if in_progress:
+                        logger.warning(f"Aborting {chat.id=}")
+                        chat_repository.set_chat_state(chat.id, interface.MessageRequestState.aborted_by_worker)
+                    else:
+                        logger.warning(f"Marking {chat.id=} as pending since no work was done.")
+                        chat_repository.set_chat_state(chat.id, interface.MessageRequestState.pending)
+                        await work_queue.enqueue(chat.id)
+                    raise
+
+                chat_repository.set_chat_state(chat.id, interface.MessageRequestState.complete)
     except fastapi.WebSocketException:
         logger.exception("Websocket closed")