92f771d455
Signed-off-by: Ian Eaves <ian.k.eaves@gmail.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Samit <285365963@qq.com> Co-authored-by: Gao Han <gaohan19@huawei.com>
75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import asyncio
|
|
from typing import Any, Generic, TypeVar
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from vllm_omni.entrypoints.openai.protocol.videos import VideoResponse
|
|
|
|
T = TypeVar("T", bound=BaseModel)
|
|
U = TypeVar("U")
|
|
|
|
|
|
class TaskRegistry:
|
|
def __init__(self) -> None:
|
|
self._tasks: dict[str, asyncio.Task[None]] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def get(self, key: str) -> asyncio.Task[None] | None:
|
|
async with self._lock:
|
|
return self._tasks.get(key)
|
|
|
|
async def pop(self, key: str) -> asyncio.Task[None] | None:
|
|
async with self._lock:
|
|
return self._tasks.pop(key, None)
|
|
|
|
async def upsert(self, key: str, task: asyncio.Task[None]) -> None:
|
|
def _cleanup(_: asyncio.Task[None]) -> None:
|
|
asyncio.create_task(self.pop(key))
|
|
|
|
task.add_done_callback(_cleanup)
|
|
|
|
async with self._lock:
|
|
self._tasks[key] = task
|
|
|
|
|
|
class AsyncDictStore(Generic[T]):
|
|
"""A small async-safe in-memory key-value store for dict items.
|
|
|
|
This encapsulates the usual pattern of a module-level dict guarded by
|
|
an asyncio.Lock and provides simple CRUD methods that are safe to call
|
|
concurrently from FastAPI request handlers and background tasks.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._items: dict[str, T] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def upsert(self, key: str, value: T) -> None:
|
|
async with self._lock:
|
|
self._items[key] = value
|
|
|
|
async def update_fields(self, key: str, updates: dict[str, Any]) -> T | None:
|
|
async with self._lock:
|
|
item: T | None = self._items.get(key)
|
|
if item is None:
|
|
return None
|
|
new_item = item.model_copy(update=updates)
|
|
self._items[key] = new_item
|
|
return new_item
|
|
|
|
async def get(self, key: str) -> T | None:
|
|
async with self._lock:
|
|
return self._items.get(key)
|
|
|
|
async def pop(self, key: str) -> T | None:
|
|
async with self._lock:
|
|
return self._items.pop(key, None)
|
|
|
|
async def list_values(self) -> list[T]:
|
|
async with self._lock:
|
|
return list(self._items.values())
|
|
|
|
|
|
VIDEO_STORE: AsyncDictStore[VideoResponse] = AsyncDictStore()
|
|
VIDEO_TASKS: TaskRegistry = TaskRegistry()
|