memryx: fix model download bug when using multiple detectors (#20030)

* Add locking for model download files

* ruff format

---------

Co-authored-by: Abinila Siva <abinila.siva@memryx.com>
This commit is contained in:
Tim Wesley
2025-09-15 09:48:55 -04:00
committed by GitHub
parent 03fe054078
commit 6cd1d1f205

View File

@@ -177,6 +177,29 @@ class MemryXDetector(DetectionApi):
logger.error(f"Failed to initialize MemryX model: {e}") logger.error(f"Failed to initialize MemryX model: {e}")
raise raise
def _acquire_file_lock(self, lock_path: str, timeout: int = 60, poll: float = 0.2):
"""
Create an exclusive lock file. Blocks (with polling) until it can acquire,
or raises TimeoutError. Uses only stdlib (os.O_EXCL).
"""
start = time.time()
while True:
try:
fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_RDWR)
os.close(fd)
return
except FileExistsError:
if time.time() - start > timeout:
raise TimeoutError(f"Timeout waiting for lock: {lock_path}")
time.sleep(poll)
def _release_file_lock(self, lock_path: str):
"""Best-effort removal of the lock file."""
try:
os.remove(lock_path)
except FileNotFoundError:
pass
def load_yolo_constants(self): def load_yolo_constants(self):
base = f"{self.cache_dir}/{self.model_folder}" base = f"{self.cache_dir}/{self.model_folder}"
# constants for yolov9 post-processing # constants for yolov9 post-processing
@@ -188,6 +211,10 @@ class MemryXDetector(DetectionApi):
if not os.path.exists(self.cache_dir): if not os.path.exists(self.cache_dir):
os.makedirs(self.cache_dir, exist_ok=True) os.makedirs(self.cache_dir, exist_ok=True)
lock_path = os.path.join(self.cache_dir, f".{self.model_folder}.lock")
self._acquire_file_lock(lock_path)
try:
# ---------- CASE 1: user provided a custom model path ---------- # ---------- CASE 1: user provided a custom model path ----------
if self.memx_model_path: if self.memx_model_path:
if not self.memx_model_path.endswith(".zip"): if not self.memx_model_path.endswith(".zip"):
@@ -245,7 +272,9 @@ class MemryXDetector(DetectionApi):
self.memx_post_model = None self.memx_post_model = None
else: else:
# Future model types can optionally use post if present # Future model types can optionally use post if present
self.memx_post_model = post_candidates[0] if post_candidates else None self.memx_post_model = (
post_candidates[0] if post_candidates else None
)
logger.info(f"Using custom model: {self.memx_model_path}") logger.info(f"Using custom model: {self.memx_model_path}")
return return
@@ -287,7 +316,9 @@ class MemryXDetector(DetectionApi):
logger.info(f"Model extracted to {self.cache_dir}.") logger.info(f"Model extracted to {self.cache_dir}.")
# Re-assign model paths after extraction # Re-assign model paths after extraction
self.memx_model_path = os.path.join(model_subdir, self.expected_dfp_model) self.memx_model_path = os.path.join(
model_subdir, self.expected_dfp_model
)
self.memx_post_model = ( self.memx_post_model = (
os.path.join(model_subdir, self.expected_post_model) os.path.join(model_subdir, self.expected_post_model)
if self.expected_post_model if self.expected_post_model
@@ -303,7 +334,12 @@ class MemryXDetector(DetectionApi):
os.remove(zip_path) os.remove(zip_path)
logger.info("Cleaned up ZIP file after extraction.") logger.info("Cleaned up ZIP file after extraction.")
except Exception as e: except Exception as e:
logger.warning(f"Failed to remove downloaded zip {zip_path}: {e}") logger.warning(
f"Failed to remove downloaded zip {zip_path}: {e}"
)
finally:
self._release_file_lock(lock_path)
def send_input(self, connection_id, tensor_input: np.ndarray): def send_input(self, connection_id, tensor_input: np.ndarray):
"""Pre-process (if needed) and send frame to MemryX input queue""" """Pre-process (if needed) and send frame to MemryX input queue"""