python break process if memory consumption is to high

Question

I am using a module and sometimes it crashes, because of memory consumption. It is terminated by SIGKILL 9 and the script is interrupted. However I cannot access the loop in which the memory consumption happens.

This is how it looks now:

import module
output = module.function(args)

and this sort of how it should look like:

import module
if ramconsumption(module.function(args)) > 999:
     output = None
else:
     output = module.function(args)

Do you know a way on how to implement this? My example is just for better understanding, the solution should just be one where I do not get a SIGKILL, when too much memory is consumed.

Answer 1

This might work:

import os
import psutil
import module
import threading

def checkMemory():
    while True:
        process = psutil.Process(os.getpid())
        memoryUsage = process.memory_info().rss #in bytes
        if memoryUsage > amount:
            print("Memory budget exceeded")
            break
    os._exit(0)

threading.Thread(name="Memory Regulator", target=checkMemory).start()

output = module.function(args)