import os
import psutil
import resource
from psutil import virtual_memory
[docs]
class RAM(object):
def __init__(self, unit="GB"):
""" helper util to manage RAM usage of this process
:param unit: string as coded in self.units
"""
self.units = {"GB": 1024 ** 3}
self.unit = unit
self.to_byte = self.units[self.unit]
[docs]
def used(self):
""" currently used ram by process """
return psutil.Process(os.getpid()).memory_info().rss / self.to_byte
[docs]
def get_limit(self):
""" limit ram usage of process, either by rlimit or hardware memory """
lim = resource.getrlimit(resource.RLIMIT_AS)[0]
if lim == -1: # unlimited
lim = virtual_memory().total
return lim / self.to_byte
[docs]
def free(self):
""" return free memory """
return self.get_limit() - self.used()
[docs]
def set_limit(self, limit):
""" set soft limit for ram usage
:param limit:
"""
if limit is not None:
resource.setrlimit(resource.RLIMIT_AS, (limit * self.to_byte, resource.RLIM_INFINITY))