mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 23:36:31 +00:00
Alex Bezdieniezhnykh
105 lines
3.8 KiB
Cython
105 lines
3.8 KiB
Cython
import os
|
|
import platform
|
|
import tempfile
|
|
from pathlib import Path
|
|
from libc.stdio cimport FILE, fopen, fclose, remove
|
|
from libc.stdlib cimport free
|
|
from libc.string cimport strdup
|
|
|
|
cdef class SecureModelLoader:
|
|
def __cinit__(self, int disk_size_mb=512):
|
|
self._disk_size_mb = disk_size_mb
|
|
self._ramdisk_path = None
|
|
self._temp_file_path = None
|
|
|
|
cpdef str load_model(self, bytes model_bytes):
|
|
"""Public method to load YOLO model securely."""
|
|
self._model_bytes = model_bytes
|
|
self._create_ramdisk()
|
|
self._store_model()
|
|
return self._temp_file_path
|
|
|
|
cdef str _get_ramdisk_path(self):
|
|
"""Determine the RAM disk path based on the OS."""
|
|
if platform.system() == "Windows":
|
|
return "R:\\"
|
|
elif platform.system() == "Linux":
|
|
return "/mnt/ramdisk"
|
|
elif platform.system() == "Darwin":
|
|
return "/Volumes/RAMDisk"
|
|
else:
|
|
raise RuntimeError("Unsupported OS for RAM disk")
|
|
|
|
cdef void _create_ramdisk(self):
|
|
"""Create a RAM disk securely based on the OS."""
|
|
system = platform.system()
|
|
|
|
if system == "Windows":
|
|
# Create RAM disk via PowerShell
|
|
command = f'powershell -Command "subst R: {tempfile.gettempdir()}"'
|
|
if os.system(command) != 0:
|
|
raise RuntimeError("Failed to create RAM disk on Windows")
|
|
self._ramdisk_path = "R:\\"
|
|
|
|
elif system == "Linux":
|
|
# Use tmpfs for RAM disk
|
|
self._ramdisk_path = "/mnt/ramdisk"
|
|
if not Path(self._ramdisk_path).exists():
|
|
os.mkdir(self._ramdisk_path)
|
|
if os.system(f"mount -t tmpfs -o size={self._disk_size_mb}M tmpfs {self._ramdisk_path}") != 0:
|
|
raise RuntimeError("Failed to create RAM disk on Linux")
|
|
|
|
elif system == "Darwin":
|
|
# Use hdiutil for macOS RAM disk
|
|
block_size = 2048 # 512-byte blocks * 2048 = 1MB
|
|
num_blocks = self._disk_size_mb * block_size
|
|
result = os.popen(f"hdiutil attach -nomount ram://{num_blocks}").read().strip()
|
|
if result:
|
|
self._ramdisk_path = "/Volumes/RAMDisk"
|
|
os.system(f"diskutil eraseVolume HFS+ RAMDisk {result}")
|
|
else:
|
|
raise RuntimeError("Failed to create RAM disk on macOS")
|
|
|
|
cdef void _store_model(self):
|
|
"""Write model securely to the RAM disk."""
|
|
cdef char* temp_path
|
|
cdef FILE* cfile
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
dir=self._ramdisk_path, suffix='.pt', delete=False
|
|
) as tmp_file:
|
|
tmp_file.write(self._model_bytes)
|
|
self._temp_file_path = tmp_file.name
|
|
|
|
encoded_path = self._temp_file_path.encode('utf-8')
|
|
temp_path = strdup(encoded_path)
|
|
with nogil:
|
|
cfile = fopen(temp_path, "rb")
|
|
if cfile == NULL:
|
|
raise IOError(f"Could not open {self._temp_file_path}")
|
|
fclose(cfile)
|
|
|
|
cdef void _cleanup(self):
|
|
"""Remove the model file and unmount RAM disk securely."""
|
|
cdef char* c_path
|
|
if self._temp_file_path:
|
|
c_path = strdup(os.fsencode(self._temp_file_path))
|
|
with nogil:
|
|
remove(c_path)
|
|
free(c_path)
|
|
self._temp_file_path = None
|
|
|
|
# Unmount RAM disk based on OS
|
|
if self._ramdisk_path:
|
|
if platform.system() == "Windows":
|
|
os.system("subst R: /D")
|
|
elif platform.system() == "Linux":
|
|
os.system(f"umount {self._ramdisk_path}")
|
|
elif platform.system() == "Darwin":
|
|
os.system("hdiutil detach /Volumes/RAMDisk")
|
|
self._ramdisk_path = None
|
|
|
|
def __dealloc__(self):
|
|
"""Ensure cleanup when the object is deleted."""
|
|
self._cleanup()
|