91 lines
2.1 KiB
Python
91 lines
2.1 KiB
Python
"""
|
|
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
|