fix api completion_lock (#6)

This commit is contained in:
josc146 2023-05-24 11:45:55 +08:00
parent 41142f15fb
commit c741b2a203

View File

@ -90,8 +90,9 @@ async def chat_completions(body: ChatCompletionBody, request: Request):
], ],
} }
) )
if await request.is_disconnected(): # torch_gc()
completion_lock.release() completion_lock.release()
if await request.is_disconnected():
return return
yield json.dumps( yield json.dumps(
{ {
@ -116,8 +117,9 @@ async def chat_completions(body: ChatCompletionBody, request: Request):
): ):
if await request.is_disconnected(): if await request.is_disconnected():
break break
if await request.is_disconnected(): # torch_gc()
completion_lock.release() completion_lock.release()
if await request.is_disconnected():
return return
yield { yield {
"response": response, "response": response,
@ -133,8 +135,6 @@ async def chat_completions(body: ChatCompletionBody, request: Request):
} }
], ],
} }
# torch_gc()
completion_lock.release()
if body.stream: if body.stream:
return EventSourceResponse(eval_rwkv()) return EventSourceResponse(eval_rwkv())
@ -182,8 +182,9 @@ async def completions(body: CompletionBody, request: Request):
], ],
} }
) )
if await request.is_disconnected(): # torch_gc()
completion_lock.release() completion_lock.release()
if await request.is_disconnected():
return return
yield json.dumps( yield json.dumps(
{ {
@ -206,8 +207,9 @@ async def completions(body: CompletionBody, request: Request):
): ):
if await request.is_disconnected(): if await request.is_disconnected():
break break
if await request.is_disconnected(): # torch_gc()
completion_lock.release() completion_lock.release()
if await request.is_disconnected():
return return
yield { yield {
"response": response, "response": response,
@ -220,8 +222,6 @@ async def completions(body: CompletionBody, request: Request):
} }
], ],
} }
# torch_gc()
completion_lock.release()
if body.stream: if body.stream:
return EventSourceResponse(eval_rwkv()) return EventSourceResponse(eval_rwkv())