43cd7f57e8
Add a sequential lock and wait until jobs are completed before executing any loading requests that directly alter the model. However, we also need to block any new requests that come in until the load is finished, so add a condition that triggers once the lock is free. Signed-off-by: kingbri <bdashore3@proton.me>
32 lines
877 B
Python
32 lines
877 B
Python
"""Concurrency handling"""
|
|
|
|
import asyncio
|
|
from fastapi.concurrency import run_in_threadpool # noqa
|
|
from typing import AsyncGenerator, Generator
|
|
|
|
|
|
# Originally from https://github.com/encode/starlette/blob/master/starlette/concurrency.py
|
|
# Uses generators instead of generics
|
|
class _StopIteration(Exception):
|
|
"""Wrapper for StopIteration because it doesn't send across threads."""
|
|
|
|
pass
|
|
|
|
|
|
def gen_next(generator: Generator):
|
|
"""Threaded function to get the next value in an iterator."""
|
|
|
|
try:
|
|
return next(generator)
|
|
except StopIteration as e:
|
|
raise _StopIteration from e
|
|
|
|
|
|
async def iterate_in_threadpool(generator: Generator) -> AsyncGenerator:
|
|
"""Iterates a generator within a threadpool."""
|
|
|
|
while True:
|
|
try:
|
|
yield await asyncio.to_thread(gen_next, generator)
|
|
except _StopIteration:
|
|
break
|