Make all SHM management untracked (blakeblackshear#15011)

ardata-tech · Nov 15, 2024 · 206ed06 · 206ed06
1 parent e407ba4
commit 206ed06
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 16 deletions.
diff --git a/frigate/app.py b/frigate/app.py
@@ -68,6 +68,7 @@
 from frigate.storage import StorageMaintainer
 from frigate.timeline import TimelineProcessor
 from frigate.util.builtin import empty_and_close_queue
+from frigate.util.image import UntrackedSharedMemory
 from frigate.util.object import get_camera_regions_grid
 from frigate.version import VERSION
 from frigate.video import capture_camera, track_camera
@@ -325,20 +326,20 @@ def start_detectors(self) -> None:
                         for det in self.config.detectors.values()
                     ]
                 )
-                shm_in = mp.shared_memory.SharedMemory(
+                shm_in = UntrackedSharedMemory(
                     name=name,
                     create=True,
                     size=largest_frame,
                 )
             except FileExistsError:
-                shm_in = mp.shared_memory.SharedMemory(name=name)
+                shm_in = UntrackedSharedMemory(name=name)
 
             try:
-                shm_out = mp.shared_memory.SharedMemory(
+                shm_out = UntrackedSharedMemory(
                     name=f"out-{name}", create=True, size=20 * 6 * 4
                 )
             except FileExistsError:
-                shm_out = mp.shared_memory.SharedMemory(name=f"out-{name}")
+                shm_out = UntrackedSharedMemory(name=f"out-{name}")
 
             self.detection_shms.append(shm_in)
             self.detection_shms.append(shm_out)

diff --git a/frigate/object_detection.py b/frigate/object_detection.py
@@ -19,7 +19,7 @@
 )
 from frigate.detectors.plugins.rocm import DETECTOR_KEY as ROCM_DETECTOR_KEY
 from frigate.util.builtin import EventsPerSecond, load_labels
-from frigate.util.image import SharedMemoryFrameManager
+from frigate.util.image import SharedMemoryFrameManager, UntrackedSharedMemory
 from frigate.util.services import listen
 
 logger = logging.getLogger(__name__)
@@ -122,7 +122,7 @@ def receiveSignal(signalNumber, frame):
 
     outputs = {}
     for name in out_events.keys():
-        out_shm = mp.shared_memory.SharedMemory(name=f"out-{name}", create=False)
+        out_shm = UntrackedSharedMemory(name=f"out-{name}", create=False)
         out_np = np.ndarray((20, 6), dtype=np.float32, buffer=out_shm.buf)
         outputs[name] = {"shm": out_shm, "np": out_np}
 
@@ -212,15 +212,13 @@ def __init__(self, name, labels, detection_queue, event, model_config, stop_even
         self.detection_queue = detection_queue
         self.event = event
         self.stop_event = stop_event
-        self.shm = mp.shared_memory.SharedMemory(name=self.name, create=False)
+        self.shm = UntrackedSharedMemory(name=self.name, create=False)
         self.np_shm = np.ndarray(
             (1, model_config.height, model_config.width, 3),
             dtype=np.uint8,
             buffer=self.shm.buf,
         )
-        self.out_shm = mp.shared_memory.SharedMemory(
-            name=f"out-{self.name}", create=False
-        )
+        self.out_shm = UntrackedSharedMemory(name=f"out-{self.name}", create=False)
         self.out_np_shm = np.ndarray((20, 6), dtype=np.float32, buffer=self.out_shm.buf)
 
     def detect(self, tensor_input, threshold=0.4):

diff --git a/frigate/util/image.py b/frigate/util/image.py
@@ -733,7 +733,7 @@ def delete(self, name):
         pass
 
 
-class SharedMemory(_mpshm.SharedMemory):
+class UntrackedSharedMemory(_mpshm.SharedMemory):
     # https://github.com/python/cpython/issues/82300#issuecomment-2169035092
 
     __lock = threading.Lock()
@@ -744,7 +744,7 @@ def __init__(
         create: bool = False,
         size: int = 0,
         *,
-        track: bool = True,
+        track: bool = False,
     ) -> None:
         self._track = track
 
@@ -779,10 +779,14 @@ def unlink(self) -> None:
 
 class SharedMemoryFrameManager(FrameManager):
     def __init__(self):
-        self.shm_store: dict[str, SharedMemory] = {}
+        self.shm_store: dict[str, UntrackedSharedMemory] = {}
 
     def create(self, name: str, size) -> AnyStr:
-        shm = SharedMemory(name=name, create=True, size=size, track=False)
+        shm = UntrackedSharedMemory(
+            name=name,
+            create=True,
+            size=size,
+        )
         self.shm_store[name] = shm
         return shm.buf
 
@@ -791,7 +795,7 @@ def get(self, name: str, shape) -> Optional[np.ndarray]:
             if name in self.shm_store:
                 shm = self.shm_store[name]
             else:
-                shm = SharedMemory(name=name, track=False)
+                shm = UntrackedSharedMemory(name=name)
                 self.shm_store[name] = shm
             return np.ndarray(shape, dtype=np.uint8, buffer=shm.buf)
         except FileNotFoundError:
@@ -814,7 +818,7 @@ def delete(self, name: str):
             del self.shm_store[name]
         else:
             try:
-                shm = SharedMemory(name=name, track=False)
+                shm = UntrackedSharedMemory(name=name)
                 shm.close()
                 shm.unlink()
             except FileNotFoundError: