161 lines
4 KiB
Python
161 lines
4 KiB
Python
"""
|
|
FileModelStorage - Local filesystem implementation of IModelStorage.
|
|
|
|
Stores model files in a local directory structure.
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
import logging
|
|
|
|
from .interface import IModelStorage
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FileModelStorage(IModelStorage):
|
|
"""
|
|
Local filesystem storage for model files.
|
|
|
|
Stores files in a directory structure:
|
|
./models/trtptcache/
|
|
trt/
|
|
<hash1>.trt
|
|
<hash2>.trt
|
|
pt/
|
|
<hash3>.pt
|
|
"""
|
|
|
|
def __init__(self, base_path: str = "./models/trtptcache"):
|
|
"""
|
|
Initialize file storage.
|
|
|
|
Args:
|
|
base_path: Base directory for storing files (default: ./models/trtptcache)
|
|
"""
|
|
self.base_path = Path(base_path).resolve()
|
|
self._ensure_directories()
|
|
|
|
def _ensure_directories(self):
|
|
"""Create base directory structure if it doesn't exist"""
|
|
self.base_path.mkdir(parents=True, exist_ok=True)
|
|
logger.info(f"Model storage initialized at: {self.base_path}")
|
|
|
|
def _get_full_path(self, key: str) -> Path:
|
|
"""
|
|
Get full filesystem path for a key.
|
|
|
|
Args:
|
|
key: Storage key (e.g., "trt/hash123.trt")
|
|
|
|
Returns:
|
|
Full filesystem path
|
|
"""
|
|
return self.base_path / key
|
|
|
|
def write(self, key: str, data: bytes) -> None:
|
|
"""
|
|
Write data to filesystem.
|
|
|
|
Args:
|
|
key: Storage key (e.g., "trt/hash123.trt")
|
|
data: Binary data to write
|
|
|
|
Raises:
|
|
IOError: If write operation fails
|
|
"""
|
|
file_path = self._get_full_path(key)
|
|
|
|
# Ensure parent directory exists
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
with open(file_path, 'wb') as f:
|
|
f.write(data)
|
|
logger.debug(f"Wrote {len(data)} bytes to {file_path}")
|
|
except Exception as e:
|
|
raise IOError(f"Failed to write to {file_path}: {e}")
|
|
|
|
def read(self, key: str) -> Optional[bytes]:
|
|
"""
|
|
Read data from filesystem.
|
|
|
|
Args:
|
|
key: Storage key
|
|
|
|
Returns:
|
|
Binary data if found, None otherwise
|
|
|
|
Raises:
|
|
IOError: If read operation fails
|
|
"""
|
|
file_path = self._get_full_path(key)
|
|
|
|
if not file_path.exists():
|
|
return None
|
|
|
|
try:
|
|
with open(file_path, 'rb') as f:
|
|
data = f.read()
|
|
logger.debug(f"Read {len(data)} bytes from {file_path}")
|
|
return data
|
|
except Exception as e:
|
|
raise IOError(f"Failed to read from {file_path}: {e}")
|
|
|
|
def exists(self, key: str) -> bool:
|
|
"""
|
|
Check if file exists.
|
|
|
|
Args:
|
|
key: Storage key
|
|
|
|
Returns:
|
|
True if file exists, False otherwise
|
|
"""
|
|
return self._get_full_path(key).exists()
|
|
|
|
def delete(self, key: str) -> bool:
|
|
"""
|
|
Delete file from filesystem.
|
|
|
|
Args:
|
|
key: Storage key
|
|
|
|
Returns:
|
|
True if deleted successfully, False if file didn't exist
|
|
"""
|
|
file_path = self._get_full_path(key)
|
|
|
|
if not file_path.exists():
|
|
return False
|
|
|
|
try:
|
|
file_path.unlink()
|
|
logger.debug(f"Deleted {file_path}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete {file_path}: {e}")
|
|
return False
|
|
|
|
def get_local_path(self, key: str) -> Optional[str]:
|
|
"""
|
|
Get local filesystem path for a key.
|
|
|
|
Args:
|
|
key: Storage key
|
|
|
|
Returns:
|
|
Local path if file exists, None otherwise
|
|
"""
|
|
file_path = self._get_full_path(key)
|
|
return str(file_path) if file_path.exists() else None
|
|
|
|
def get_storage_path(self) -> str:
|
|
"""
|
|
Get the base storage path.
|
|
|
|
Returns:
|
|
Base path where files are stored
|
|
"""
|
|
return str(self.base_path)
|