wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-08-12 11:40:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

robustifying inference

Browse Source
This commit is contained in:
Yannic Kilcher
2023-02-09 15:31:46 +01:00
parent 7c4ff73241
commit ed7d920e5d
2 changed files with 32 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
inference/server/main.py
+8 -1
View File
@@ -216,17 +216,24 @@ async def work(websocket: fastapi.WebSocket):
break
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:
break
except fastapi.WebSocketException:
# TODO: handle this better
logger.exception(f"Websocket closed during handling of {chat.id}")
chat.message_request_state = MessageRequestState.aborted_by_worker
if in_progress:
logger.warning(f"Aborting {chat.id=}")
chat.message_request_state = MessageRequestState.aborted_by_worker
else:
logger.warning(f"Marking {chat.id=} as pending since no work was done.")
chat.message_request_state = MessageRequestState.pending
raise
chat.message_request_state = MessageRequestState.complete