Last active
October 19, 2023 09:17
-
-
Save eladn/10d9374dabfa7cb4e313a79fe207d5e5 to your computer and use it in GitHub Desktop.
Python memory watchdog - Monitors a process' memory consumption; Gracefully triggers an exception in the monitored process if exceeds memory usage soft limit or kills the process if exceeds the hard limit
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| __author__ = "Elad Nachmias" | |
| __email__ = "eladnah@gmail.com" | |
| __date__ = "2023-10-18" | |
| import multiprocessing | |
| import os | |
| import signal | |
| import time | |
| import psutil | |
| import warnings | |
| SoftMaxAllowedMemoryConsumptionExceededSignal = signal.SIGUSR2 | |
| class MemoryWatchdog(multiprocessing.Process): | |
| """ | |
| Monitoring a process' memory consumption. Gracefully triggers an exception in the monitored process if exceeds | |
| memory usage soft limit or agressively kills the process if exceeds the hard limit. | |
| """ | |
| def __init__( | |
| self, | |
| monitored_process_identifier: str, | |
| process_id_to_monitor: int, | |
| soft_max_allowed_memory_consumption: int, | |
| hard_max_allowed_memory_consumption: int): | |
| super(MemoryWatchdog, self).__init__() | |
| self._monitored_process_identifier = monitored_process_identifier | |
| self._process_id_to_monitor = process_id_to_monitor | |
| self._soft_max_allowed_memory_consumption = soft_max_allowed_memory_consumption | |
| self._hard_max_allowed_memory_consumption = hard_max_allowed_memory_consumption | |
| @property | |
| def _process_to_monitor(self): | |
| return psutil.Process(pid=self._process_id_to_monitor) | |
| def run(self) -> None: | |
| while psutil.pid_exists(self._process_id_to_monitor) and self._process_to_monitor.is_running(): | |
| memory_consumption = self._process_to_monitor.memory_info().rss | |
| if memory_consumption > self._hard_max_allowed_memory_consumption: | |
| warnings.warn( | |
| f'Killing process {self._monitored_process_identifier} [{self._process_id_to_monitor}] because of exceeding ' | |
| f'of hard memory consumption limit ({memory_consumption} > {self._hard_max_allowed_memory_consumption})') | |
| os.kill(self._process_id_to_monitor, signal.SIGKILL) | |
| elif memory_consumption > self._soft_max_allowed_memory_consumption: | |
| # Send some custom user signal, register on it from the worker's process, and exit gracefully from within there. | |
| warnings.warn( | |
| f'Trying to handle gracefully soft memory consumption limit exceeding @ process ' | |
| f'{self._monitored_process_identifier} [{self._process_id_to_monitor}] ' | |
| f'({memory_consumption} > {self._soft_max_allowed_memory_consumption})') | |
| os.kill(self._process_id_to_monitor, SoftMaxAllowedMemoryConsumptionExceededSignal) | |
| time.sleep(10.) | |
| def register(self): | |
| signal.signal(SoftMaxAllowedMemoryConsumptionExceededSignal, _soft_max_allowed_memory_consumption_exceeded_sig_handler) | |
| class SoftMaxAllowedMemoryConsumptionExceededError(Exception): | |
| pass | |
| def _soft_max_allowed_memory_consumption_exceeded_sig_handler(signum, frame): | |
| raise SoftMaxAllowedMemoryConsumptionExceededError |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment