|
10 | 10 | from collections.abc import Callable |
11 | 11 | from dataclasses import dataclass |
12 | 12 | from multiprocessing import Queue, shared_memory |
13 | | -from multiprocessing.synchronize import Event, Semaphore |
| 13 | +from multiprocessing.synchronize import Event |
14 | 14 | from queue import Empty |
15 | 15 |
|
16 | 16 | import numpy as np |
17 | 17 | import psutil |
18 | 18 |
|
19 | 19 | import birdnet.acoustic.inference.core.logs as bn_logging |
20 | 20 | from birdnet.acoustic.inference.core.shm import RingField |
| 21 | +from birdnet.acoustic.inference.core.sync import CountedSemaphore |
21 | 22 | from birdnet.globals import READABLE_FLAG, READING_FLAG, WRITABLE_FLAG |
22 | 23 |
|
23 | 24 |
|
@@ -124,8 +125,8 @@ def __init__( |
124 | 125 | logging_queue: Queue, |
125 | 126 | logging_level: int, |
126 | 127 | perf_res: Queue, |
127 | | - sem_active_workers: Semaphore, |
128 | | - sem_filled_slots: Semaphore, |
| 128 | + sem_active_workers: CountedSemaphore, |
| 129 | + sem_filled_slots: CountedSemaphore, |
129 | 130 | segment_size_s: float, |
130 | 131 | parent_process_id: int, |
131 | 132 | rf_flags: RingField, |
@@ -275,20 +276,42 @@ def reset(self) -> None: |
275 | 276 | self._prd_speed_xrt_tracker.reset() |
276 | 277 | self._prd_speed_seg_per_s_tracker.reset() |
277 | 278 |
|
| 279 | + @staticmethod |
| 280 | + def _safe_proc_memory(proc: psutil.Process) -> float | None: |
| 281 | + try: |
| 282 | + return float(proc.memory_full_info().uss) |
| 283 | + except (psutil.AccessDenied, PermissionError): |
| 284 | + pass |
| 285 | + except psutil.NoSuchProcess: |
| 286 | + return None |
| 287 | + try: |
| 288 | + return float(proc.memory_info().rss) |
| 289 | + except (psutil.NoSuchProcess, psutil.AccessDenied, PermissionError): |
| 290 | + return None |
| 291 | + |
278 | 292 | def _track_memory_usage(self) -> None: |
279 | 293 | if self._parent_process is None: |
280 | | - self._parent_process = psutil.Process(self._parent_process_id) |
281 | | - memory_usage: float = self._parent_process.memory_full_info().uss |
282 | | - for child in self._parent_process.children(recursive=True): |
283 | 294 | try: |
284 | | - memory_usage += child.memory_full_info().uss |
285 | | - except psutil.NoSuchProcess: |
286 | | - continue |
287 | | - except psutil.AccessDenied: |
288 | | - continue |
289 | | - |
290 | | - mem_usage_MiB = memory_usage / 1024**2 |
291 | | - self._memory_usage_MiB_tracker.add_value(mem_usage_MiB) |
| 295 | + self._parent_process = psutil.Process(self._parent_process_id) |
| 296 | + except (psutil.NoSuchProcess, psutil.AccessDenied, PermissionError): |
| 297 | + return |
| 298 | + |
| 299 | + parent_mem = self._safe_proc_memory(self._parent_process) |
| 300 | + if parent_mem is None: |
| 301 | + return |
| 302 | + |
| 303 | + total = parent_mem |
| 304 | + try: |
| 305 | + children = self._parent_process.children(recursive=True) |
| 306 | + except (psutil.AccessDenied, PermissionError, psutil.NoSuchProcess): |
| 307 | + children = [] |
| 308 | + |
| 309 | + for child in children: |
| 310 | + child_mem = self._safe_proc_memory(child) |
| 311 | + if child_mem is not None: |
| 312 | + total += child_mem |
| 313 | + |
| 314 | + self._memory_usage_MiB_tracker.add_value(total / 1024**2) |
292 | 315 |
|
293 | 316 | @property |
294 | 317 | def wall_time(self) -> float: |
|
0 commit comments