theluyuan/RWKV-Runner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix api completion_lock (#6)

Browse Source
This commit is contained in:
josc146 2023-05-24 11:45:55 +08:00
parent 41142f15fb
commit c741b2a203
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
backend-python/routes/completion.py
View File

@ -90,8 +90,9 @@ async def chat_completions(body: ChatCompletionBody, request: Request):
], ],
} }
) )
# torch_gc()
completion_lock.release()
if await request.is_disconnected(): if await request.is_disconnected():
completion_lock.release()
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
# torch_gc()
completion_lock.release()
if await request.is_disconnected(): if await request.is_disconnected():
completion_lock.release()
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):
], ],
} }
) )
# torch_gc()
completion_lock.release()
if await request.is_disconnected(): if await request.is_disconnected():
completion_lock.release()
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
# torch_gc()
completion_lock.release()
if await request.is_disconnected(): if await request.is_disconnected():
completion_lock.release()
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())