converter system
This commit is contained in:
parent
d3dbf9a580
commit
748fb71980
9 changed files with 1012 additions and 14 deletions
8
services/modelstorage/__init__.py
Normal file
8
services/modelstorage/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
"""
|
||||
Model storage module for managing TensorRT and PyTorch model files.
|
||||
"""
|
||||
|
||||
from .interface import IModelStorage
|
||||
from .file_storage import FileModelStorage
|
||||
|
||||
__all__ = ['IModelStorage', 'FileModelStorage']
|
||||
161
services/modelstorage/file_storage.py
Normal file
161
services/modelstorage/file_storage.py
Normal file
|
|
@ -0,0 +1,161 @@
|
|||
"""
|
||||
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)
|
||||
91
services/modelstorage/interface.py
Normal file
91
services/modelstorage/interface.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
"""
|
||||
IModelStorage - Interface for model file storage.
|
||||
|
||||
Defines the contract for storing and retrieving model files (TensorRT, PyTorch, etc.)
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, BinaryIO
|
||||
from pathlib import Path
|
||||
import io
|
||||
|
||||
|
||||
class IModelStorage(ABC):
|
||||
"""
|
||||
Interface for model file storage.
|
||||
|
||||
This abstraction allows swapping storage backends (local filesystem, S3, etc.)
|
||||
without changing the model conversion and loading logic.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def write(self, key: str, data: bytes) -> None:
|
||||
"""
|
||||
Write data to storage with the given key.
|
||||
|
||||
Args:
|
||||
key: Storage key (e.g., "trt/hash123.trt" or "pt/hash456.pt")
|
||||
data: Binary data to write
|
||||
|
||||
Raises:
|
||||
IOError: If write operation fails
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def read(self, key: str) -> Optional[bytes]:
|
||||
"""
|
||||
Read data from storage by key.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
Binary data if found, None otherwise
|
||||
|
||||
Raises:
|
||||
IOError: If read operation fails
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Check if a key exists in storage.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
True if key exists, False otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, key: str) -> bool:
|
||||
"""
|
||||
Delete data from storage.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
True if deleted successfully, False if key didn't exist
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_local_path(self, key: str) -> Optional[str]:
|
||||
"""
|
||||
Get a local filesystem path for the key.
|
||||
|
||||
For local storage, this returns the direct path.
|
||||
For remote storage (S3), this may download to a temp location or return None.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
Local path if available/downloaded, None if not supported
|
||||
"""
|
||||
pass
|
||||
Loading…
Add table
Add a link
Reference in a new issue