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>
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import asyncio
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from tempfile import NamedTemporaryFile
|
|
|
|
|
|
class StorageBaseManager(ABC):
|
|
@abstractmethod
|
|
async def save(self, *args, **kwargs) -> str:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def delete(self, *args, **kwargs) -> bool:
|
|
pass
|
|
|
|
|
|
class LocalStorageManager(StorageBaseManager):
|
|
def __init__(self, storage_path: str | None = None):
|
|
if storage_path is None:
|
|
storage_path = os.getenv("VLLM_OMNI_STORAGE_PATH", "/tmp/storage")
|
|
self.storage_path = storage_path
|
|
os.makedirs(self.storage_path, exist_ok=True)
|
|
|
|
max_concurrency = int(os.getenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "4"))
|
|
self._io_semaphore = asyncio.Semaphore(max(1, max_concurrency))
|
|
|
|
def _save_sync(self, data: bytes, file_name: str) -> str:
|
|
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)
|
|
return filename
|
|
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) -> str:
|
|
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:
|
|
return os.path.join(self.storage_path, file_name)
|
|
|
|
|
|
# Should implement TTL management on local file storage
|
|
# and recovery on restart in the *STORE's
|
|
STORAGE_MANAGER = LocalStorageManager()
|