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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -177,6 +177,29 @@ class MemryXDetector(DetectionApi):
logger.error(f"Failed to initialize MemryX model: {e}")
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):
base = f"{self.cache_dir}/{self.model_folder}"
# constants for yolov9 post-processing
@ -188,6 +211,10 @@ class MemryXDetector(DetectionApi):
if not os.path.exists(self.cache_dir):
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 ----------
if self.memx_model_path:
if not self.memx_model_path.endswith(".zip"):
@ -245,7 +272,9 @@ class MemryXDetector(DetectionApi):
self.memx_post_model = None
else:
# 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}")
return
@ -287,7 +316,9 @@ class MemryXDetector(DetectionApi):
logger.info(f"Model extracted to {self.cache_dir}.")
# 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 = (
os.path.join(model_subdir, self.expected_post_model)
if self.expected_post_model
@ -303,7 +334,12 @@ class MemryXDetector(DetectionApi):
os.remove(zip_path)
logger.info("Cleaned up ZIP file after extraction.")
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):
"""Pre-process (if needed) and send frame to MemryX input queue"""