File size: 5,526 Bytes
3943768 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
import os
from functools import wraps
import psutil
rlims = [psutil.RLIMIT_NOFILE if hasattr(psutil, 'RLIMIT_NOFILE') else None, psutil.RLIMIT_NPROC if hasattr(psutil, 'RLIMIT_NPROC') else None]
rlims_str = ["RLIMIT_NOFILE", "RLIMIT_NPROC"]
def rlimitproc(pp, rlim):
try:
return pp.rlimit(rlim)
except (psutil.NoSuchProcess, psutil.AccessDenied, FileNotFoundError, OSError, TypeError, AttributeError):
pass
except ValueError as e:
if 'invalid resource specified' in str(e):
print("rlimitproc exception for rlim %s: %s" % (rlim, str(e)))
else:
raise
except Exception as e:
print("rlimitproc exception: rlim %s: %s" % (rlim, str(e)))
if os.environ.get('HARD_ASSERTS'):
raise
pass
return -1, -1
def get_all_rlimit(pid=None):
if pid is None:
pid = os.getpid()
ps = psfunc(psutil.Process, pid)
result = {}
for rlim_str, rlim in zip(rlims_str, rlims):
if rlims is None:
continue
result[(rlim_str, rlim)] = rlimitproc(ps, rlim)
return result
limit_nofile = 131071
limit_nproc = 16384
def reulimit(pid=None, verbose=False):
from sys import platform
if not (platform == "linux" or platform == "linux2"):
return
if pid is None:
pid = os.getpid()
ps = psfunc(psutil.Process, pid)
ulimits_dict = get_all_rlimit()
for k, v in zip(ulimits_dict.keys(), ulimits_dict.values()):
if k[1] == psutil.RLIMIT_CORE:
continue
if verbose:
print("rlimit %s of %s" % (str(k[0]), str(v[0])))
if isinstance(v, tuple) and len(v) == 2:
newlimits = list(v)
# set soft to hard limit
if newlimits[0] != newlimits[1]:
if k[1] == psutil.RLIMIT_NOFILE:
hard_limit = newlimits[1] if newlimits[1] != -1 else limit_nofile
newlimits[0] = max(newlimits[0], min(limit_nofile, hard_limit))
elif k[1] == psutil.RLIMIT_NPROC:
hard_limit = newlimits[1] if newlimits[1] != -1 else limit_nproc
newlimits[0] = max(newlimits[0], min(limit_nproc, hard_limit))
else:
newlimits[0] = newlimits[1]
try:
ps.rlimit(k[1], limits=tuple(newlimits))
if verbose:
print("Set rlimit %s of %s -> %s" % (str(k[0]), str(v[0]), str(newlimits[0])))
except (TypeError, AttributeError, psutil.AccessDenied):
print("Could not set desired rlimit %s of %s -> %s" % (
str(k[0]), str(v[0]), str(newlimits[0])))
except (FileNotFoundError, OSError, psutil.NoSuchProcess):
pass
except Exception as e:
print("Couldn't set ulimit %s" % str(e))
if os.environ.get('HARD_ASSERTS'):
raise
return
def get_nproc_limit(pid=None):
if pid is None:
pid = os.getpid()
ps = psfunc(psutil.Process, pid)
if ps is not None:
nproc = rlimitproc(ps, psutil.RLIMIT_NPROC) # (soft, hard)
else:
nproc = (-1, -1)
nproc = list(nproc)
if nproc[0] == -1:
nproc[0] = limit_nproc
if nproc[1] == -1:
nproc[1] = limit_nproc
return tuple(nproc)
def wrap_psutil(func):
""" Decorate a function that uses psutil in case of ignorable exception
"""
@wraps(func)
def f(*args, **kwargs):
val = psfunc(func, *args, **kwargs)
return val
return f
def psfunc_list(func, *args, **kwargs):
ret = psfunc(func, *args, **kwargs)
if ret is None:
return []
else:
return ret
def psfunc(func, *args, **kwargs):
"""
Safely ask for psutil function call
psutil accesses /proc entries that can random disappear, and psutil does not have sufficient protection
for user against various errors either direct or a cascade within the package.
:param func: psutil function to use
:param args: args
:param kwargs: kwargs
:return: function return value
"""
try:
return func(*args, **kwargs)
except (psutil.NoSuchProcess, psutil.AccessDenied, FileNotFoundError, OSError, TypeError, AttributeError):
pass
except Exception as e:
if os.environ.get('HARD_ASSERTS'):
raise
def psattr(obj, attr):
"""
Safely ask for an attributes value for psutil
psutil accesses /proc entries that can random disappear, and psutil does not have sufficient protection
for user against various errors either direct or a cascade within the package.
:param obj: psutil object with attributes
:param attr: attribute name to get
:return: attribute value
"""
try:
return getattr(obj, attr)
except (psutil.NoSuchProcess, psutil.AccessDenied, FileNotFoundError, OSError, TypeError, AttributeError):
pass
except Exception as e:
if os.environ.get('HARD_ASSERTS'):
raise
def get_file_limit(pid=None):
if pid is None:
pid = os.getpid()
ps = psfunc(psutil.Process, pid)
if ps is not None:
nofile = rlimitproc(ps, psutil.RLIMIT_NOFILE) # (soft, hard)
else:
nofile = (-1, -1)
nofile = list(nofile)
if nofile[0] == -1:
nofile[0] = limit_nofile
if nofile[1] == -1:
nofile[1] = limit_nofile
return tuple(nofile)
|