converter system

services/modelstorage/file_storage.py

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