Skip to content

Commit 730171a

Browse files
fabianpfabianp
authored
Merge pull request #311 from mgbckr/master
Add capability to switch to PSS and USS for psutil
2 parents e43d78b + b9513a0 commit 730171a

3 files changed

Lines changed: 179 additions & 2 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
.idea
2+
.vscode
23
dist
34
build
45
MANIFEST

‎memory_profiler.py‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ def _repr_pretty_(self, p, cycle):
8282
p.text(u'<MemitResult : ' + msg + u'>')
8383

8484

85-
def _get_child_memory(process, meminfo_attr=None):
85+
def _get_child_memory(process, meminfo_attr=None, memory_metric=0):
8686
"""
8787
Returns a generator that yields memory for all child processes.
8888
"""
@@ -102,7 +102,11 @@ def _get_child_memory(process, meminfo_attr=None):
102102
# Loop over the child processes and yield their memory
103103
try:
104104
for child in getattr(process, children_attr)(recursive=True):
105-
yield getattr(child, meminfo_attr)()[0] / _TWO_20
105+
if isinstance(memory_metric, str):
106+
meminfo = getattr(child, meminfo_attr)()
107+
yield getattr(meminfo, memory_metric) / _TWO_20
108+
else:
109+
yield getattr(child, meminfo_attr)()[memory_metric] / _TWO_20
106110
except (psutil.NoSuchProcess, psutil.AccessDenied):
107111
# https://github.com/fabianp/memory_profiler/issues/71
108112
yield 0.0
@@ -142,6 +146,35 @@ def ps_util_tool():
142146
pass
143147
# continue and try to get this from ps
144148

149+
def _ps_util_full_tool(memory_metric):
150+
151+
# .. cross-platform but requires psutil > 4.0.0 ..
152+
process = psutil.Process(pid)
153+
try:
154+
if not hasattr(process, 'memory_full_info'):
155+
raise NotImplementedError("Backend `psutil_pss` and `psutil_uss` requires psutil > 4.0.0")
156+
157+
meminfo_attr = 'memory_full_info'
158+
meminfo = getattr(process, meminfo_attr)()
159+
160+
if not hasattr(meminfo, memory_metric):
161+
raise NotImplementedError(
162+
"Metric `{}` not available. For details, see:".format(memory_metric) +
163+
"https://psutil.readthedocs.io/en/latest/index.html?highlight=memory_info#psutil.Process.memory_full_info")
164+
mem = getattr(meminfo, memory_metric) / _TWO_20
165+
166+
if include_children:
167+
mem += sum(_get_child_memory(process, meminfo_attr, memory_metric))
168+
169+
if timestamps:
170+
return mem, time.time()
171+
else:
172+
return mem
173+
174+
except psutil.AccessDenied:
175+
pass
176+
# continue and try to get this from ps
177+
145178
def posix_tool():
146179
# .. scary stuff ..
147180
if include_children:
@@ -180,6 +213,8 @@ def posix_tool():
180213

181214
tools = {'tracemalloc': tracemalloc_tool,
182215
'psutil': ps_util_tool,
216+
'psutil_pss': lambda: _ps_util_full_tool(memory_metric="pss"),
217+
'psutil_uss': lambda: _ps_util_full_tool(memory_metric="uss"),
183218
'posix': posix_tool}
184219
return tools[backend]()
185220

@@ -1164,6 +1199,8 @@ def choose_backend(new_backend=None):
11641199
_backend = 'no_backend'
11651200
all_backends = [
11661201
('psutil', True),
1202+
('psutil_pss', True),
1203+
('psutil_uss', True),
11671204
('posix', os.name == 'posix'),
11681205
('tracemalloc', has_tracemalloc),
11691206
]

‎test/test_psutil_memory_full_info.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
from memory_profiler import memory_usage
2+
3+
# size = 50000
4+
size = 3000
5+
6+
7+
def test_simple():
8+
9+
import numpy as np
10+
11+
def func():
12+
a = np.random.random((size, size))
13+
return a
14+
15+
rss = memory_usage(proc=func, max_usage=True, backend="psutil")
16+
uss = memory_usage(proc=func, max_usage=True, backend="psutil_uss")
17+
pss = memory_usage(proc=func, max_usage=True, backend="psutil_pss")
18+
print(rss, uss, pss)
19+
20+
21+
def test_multiprocessing():
22+
23+
import numpy as np
24+
import joblib
25+
import time
26+
27+
def func():
28+
n_jobs = 4
29+
a = np.random.random((size, size))
30+
31+
def subprocess(i):
32+
time.sleep(2)
33+
return a[i,i]
34+
35+
results = joblib.Parallel(n_jobs=n_jobs)(
36+
joblib.delayed(subprocess)(i)
37+
for i in range(n_jobs))
38+
39+
return results
40+
41+
rss = memory_usage(proc=func, max_usage=True, backend="psutil", include_children=True, multiprocess=True)
42+
uss = memory_usage(proc=func, max_usage=True, backend="psutil_uss", include_children=True, multiprocess=True)
43+
pss = memory_usage(proc=func, max_usage=True, backend="psutil_pss", include_children=True, multiprocess=True)
44+
print(rss, uss, pss)
45+
46+
47+
def test_multiprocessing_write():
48+
49+
import numpy as np
50+
import joblib
51+
import time
52+
53+
def func():
54+
n_jobs = 4
55+
a = np.random.random((size, size))
56+
57+
def subprocess(i):
58+
aa = a.copy()
59+
time.sleep(2)
60+
return aa[i,i]
61+
62+
results = joblib.Parallel(n_jobs=n_jobs)(
63+
joblib.delayed(subprocess)(i)
64+
for i in range(n_jobs))
65+
66+
return results
67+
68+
rss = memory_usage(proc=func, max_usage=True, backend="psutil", include_children=True, multiprocess=True)
69+
uss = memory_usage(proc=func, max_usage=True, backend="psutil_uss", include_children=True, multiprocess=True)
70+
pss = memory_usage(proc=func, max_usage=True, backend="psutil_pss", include_children=True, multiprocess=True)
71+
print(rss, uss, pss)
72+
73+
74+
def test_multiprocessing_showcase():
75+
76+
import numpy as np
77+
import joblib
78+
import time
79+
import datetime
80+
81+
def func():
82+
83+
# n_jobs = 32
84+
# size = 25000
85+
# Creating data: 25000x25000 ... done (4.66 Gb). Starting processing: n_jobs=32 ... done (0:00:37.581291). RSS: 353024.01
86+
# Creating data: 25000x25000 ... done (4.66 Gb). Starting processing: n_jobs=32 ... done (0:00:38.867385). USS: 148608.62
87+
# Creating data: 25000x25000 ... done (4.66 Gb). Starting processing: n_jobs=32 ... done (0:00:29.049754). PSS: 169253.91
88+
89+
# n_jobs = 64
90+
# size = 10000
91+
# Creating data: 10000x10000 ... done (0.75 Gb). Starting processing: n_jobs=64 ... done (0:00:14.701243). RSS: 111362.79
92+
# Creating data: 10000x10000 ... done (0.75 Gb). Starting processing: n_jobs=64 ... done (0:00:15.020202). USS: 56108.69
93+
# Creating data: 10000x10000 ... done (0.75 Gb). Starting processing: n_jobs=64 ... done (0:00:15.072918). PSS: 54826.61
94+
95+
# Conclusion:
96+
# * RSS is overestimating like crazy (I checked the actual memory usage using htop)
97+
98+
n_jobs = 8
99+
size = 3000
100+
101+
print("Creating data: {size}x{size} ... ".format(size=size), end="")
102+
a = np.random.random((size, size))
103+
print("done ({size:.02f} Gb). ".format(size=a.size * a.itemsize / 1024**3), end="")
104+
105+
def subprocess(i):
106+
aa = a.copy()
107+
r = aa[1,1]
108+
aa = a.copy()
109+
time.sleep(10)
110+
return r
111+
112+
# r = a[1,1]
113+
# # time.sleep(10)
114+
# return r
115+
116+
pass
117+
118+
start = datetime.datetime.now()
119+
print("Starting processing: n_jobs={n_jobs} ... ".format(n_jobs=n_jobs), end="")
120+
results = joblib.Parallel(n_jobs=n_jobs)(
121+
joblib.delayed(subprocess)(i)
122+
for i in range(n_jobs))
123+
print("done ({}). ".format(datetime.datetime.now() - start), end="")
124+
125+
return results
126+
127+
rss = memory_usage(proc=func, max_usage=True, backend="psutil", include_children=True, multiprocess=True)
128+
print("RSS: {rss:.02f}".format(rss=rss))
129+
uss = memory_usage(proc=func, max_usage=True, backend="psutil_uss", include_children=True, multiprocess=True)
130+
print("USS: {uss:.02f}".format(uss=uss))
131+
pss = memory_usage(proc=func, max_usage=True, backend="psutil_pss", include_children=True, multiprocess=True)
132+
print("PSS: {pss:.02f}".format(pss=pss))
133+
134+
135+
if __name__ == "__main__":
136+
test_simple()
137+
test_multiprocessing()
138+
test_multiprocessing_write()
139+
test_multiprocessing_showcase()

0 commit comments

Comments
 (0)