converter system

This commit is contained in:
Siwat Sirichai 2025-11-09 19:54:35 +07:00
parent d3dbf9a580
commit 748fb71980
9 changed files with 1012 additions and 14 deletions

View 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']

View 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)

View 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