211 lines
7.1 KiB
Python
211 lines
7.1 KiB
Python
import asyncio
|
|
import contextlib
|
|
import os
|
|
import stat
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from tempfile import NamedTemporaryFile
|
|
from typing import Generic, Literal, TypeVar
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.config.server_settings import SERVER_SETTINGS_CONFIG, FileBackend
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class SaveContext:
|
|
key: str
|
|
created_at: int
|
|
expires_at: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BaseStorageHandle:
|
|
kind: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FileStorageHandle(BaseStorageHandle):
|
|
path: str
|
|
kind: Literal["path"] = field(default="path", init=False)
|
|
|
|
|
|
K = TypeVar("K", bound=BaseStorageHandle, covariant=True)
|
|
|
|
|
|
class StorageBaseManager(Generic[K], ABC):
|
|
@abstractmethod
|
|
async def save(self, *args, **kwargs) -> SaveContext:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def delete(self, *args, **kwargs) -> bool:
|
|
pass
|
|
|
|
async def start(self, *args, **kwargs) -> None:
|
|
pass
|
|
|
|
async def stop(self, *args, **kwargs) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def open(self, storage_key: str) -> K | None:
|
|
pass
|
|
|
|
|
|
class LocalStorageManager(StorageBaseManager[FileStorageHandle]):
|
|
def __init__(self, storage_path: str, max_concurrency: int = 4):
|
|
self.storage_path = os.path.realpath(storage_path)
|
|
os.makedirs(self.storage_path, exist_ok=True)
|
|
|
|
self._io_semaphore = asyncio.Semaphore(max(1, max_concurrency))
|
|
|
|
async def open(self, storage_key: str) -> FileStorageHandle | None:
|
|
local_file = self.get_full_file_path(storage_key)
|
|
if not os.path.exists(local_file):
|
|
return None
|
|
return FileStorageHandle(path=local_file)
|
|
|
|
def _save_sync(self, data: bytes, file_name: str) -> SaveContext:
|
|
filename = self.get_full_file_path(file_name)
|
|
tmp_name: str | None = None
|
|
|
|
try:
|
|
with NamedTemporaryFile("wb", dir=self.storage_path, delete=False) as f:
|
|
tmp_name = f.name
|
|
f.write(data)
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
os.replace(tmp_name, filename)
|
|
response = SaveContext(key=file_name, created_at=int(time.time()))
|
|
return response
|
|
except Exception:
|
|
if tmp_name is not None:
|
|
try:
|
|
os.remove(tmp_name)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
async def save(self, data: bytes, file_name: str) -> SaveContext:
|
|
async with self._io_semaphore:
|
|
return await asyncio.to_thread(self._save_sync, data, file_name)
|
|
|
|
def _delete_sync(self, file_name: str) -> bool:
|
|
try:
|
|
os.remove(self.get_full_file_path(file_name))
|
|
except FileNotFoundError:
|
|
return False
|
|
return True
|
|
|
|
async def delete(self, file_name: str) -> bool:
|
|
async with self._io_semaphore:
|
|
return await asyncio.to_thread(self._delete_sync, file_name)
|
|
|
|
def exists(self, file_name: str) -> bool:
|
|
return os.path.exists(self.get_full_file_path(file_name))
|
|
|
|
def get_full_file_path(self, file_name: str) -> str:
|
|
full = os.path.realpath(os.path.join(self.storage_path, file_name))
|
|
|
|
if os.path.commonpath([self.storage_path, full]) != self.storage_path:
|
|
raise ValueError(f"Illegal storage key: {file_name!r}")
|
|
|
|
return full
|
|
|
|
|
|
class LocalStorageTTLManager(LocalStorageManager):
|
|
def __init__(self, ttl_seconds: int, sweep_interval_seconds: int, *args, **kwargs):
|
|
if ttl_seconds <= 0:
|
|
raise ValueError("`ttl_seconds` must be greater than or equal to 1.")
|
|
if sweep_interval_seconds <= 0:
|
|
raise ValueError("`sweep_interval_seconds` must be greater than or equal to 1.")
|
|
|
|
self._ttl_seconds = ttl_seconds
|
|
self._sweep_interval_seconds = sweep_interval_seconds
|
|
self._sweeper_task: asyncio.Task[None] | None = None
|
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
async def save(self, data: bytes, file_name: str) -> SaveContext:
|
|
result = await super().save(data, file_name)
|
|
result.expires_at = result.created_at + self._ttl_seconds
|
|
return result
|
|
|
|
async def _sweep_once(self, cutoff: float) -> int:
|
|
expired: list[str] = []
|
|
for entry in os.scandir(self.storage_path):
|
|
try:
|
|
if not entry.is_file(follow_symlinks=False):
|
|
continue
|
|
if entry.stat(follow_symlinks=False).st_mtime < cutoff:
|
|
expired.append(entry.path)
|
|
except FileNotFoundError:
|
|
logger.debug("TTL sweep skipped %s; file removed during scan", entry.path)
|
|
except OSError:
|
|
logger.warning("TTL sweep failed to inspect %s", entry.path, exc_info=True)
|
|
|
|
deleted = 0
|
|
for path in expired:
|
|
try:
|
|
async with self._io_semaphore:
|
|
st = await asyncio.to_thread(os.stat, path, follow_symlinks=False)
|
|
if not stat.S_ISREG(st.st_mode):
|
|
continue
|
|
if st.st_mtime >= cutoff:
|
|
continue
|
|
await asyncio.to_thread(os.remove, path)
|
|
deleted += 1
|
|
except FileNotFoundError:
|
|
logger.debug("TTL sweep skipped %s; file already removed", path)
|
|
except OSError:
|
|
logger.warning("TTL sweep failed to delete expired file %s", path, exc_info=True)
|
|
|
|
return deleted
|
|
|
|
async def _sweep_loop(self) -> None:
|
|
while True:
|
|
try:
|
|
cutoff = time.time() - self._ttl_seconds
|
|
await self._sweep_once(cutoff)
|
|
except Exception:
|
|
logger.exception("TTL sweep failed for storage path %s", self.storage_path)
|
|
await asyncio.sleep(self._sweep_interval_seconds)
|
|
|
|
async def start(self) -> None:
|
|
if self._sweeper_task is None or self._sweeper_task.done():
|
|
self._sweeper_task = asyncio.create_task(self._sweep_loop())
|
|
|
|
async def stop(self) -> None:
|
|
if self._sweeper_task is None:
|
|
return
|
|
self._sweeper_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await self._sweeper_task
|
|
self._sweeper_task = None
|
|
|
|
|
|
def get_storage_manager(storage_config: FileBackend) -> StorageBaseManager[FileStorageHandle]:
|
|
if isinstance(storage_config, FileBackend):
|
|
if storage_config.file_ttl is not None and storage_config.ttl_sweep_interval is not None:
|
|
manager = LocalStorageTTLManager(
|
|
storage_path=storage_config.path,
|
|
max_concurrency=storage_config.file_concurrency,
|
|
ttl_seconds=storage_config.file_ttl,
|
|
sweep_interval_seconds=storage_config.ttl_sweep_interval,
|
|
)
|
|
else:
|
|
manager = LocalStorageManager(
|
|
storage_path=storage_config.path, max_concurrency=storage_config.file_concurrency
|
|
)
|
|
else:
|
|
raise ValueError("No supported storage managers")
|
|
|
|
return manager
|
|
|
|
|
|
STORAGE_MANAGER = get_storage_manager(SERVER_SETTINGS_CONFIG.storage)
|