/
process.py
173 lines (146 loc) · 4.87 KB
/
process.py
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
"""Tools to read out total memory usage and get/set the number of threads.
If your python is compiled against MKL (e.g. if you use *anaconda* as recommended in INSTALL),
it will by default use as many threads as CPU cores are available.
If you run a job on a cluster, you should limit this to the number of cores you reserved --
otherwise your colleagues might get angry...
A simple way to achieve this is to set a suitable environment variable before calling your
python program, e.g. on the linux bash ``export OMP_NUM_THREADS=4`` for 4 threads.
(MKL used OpenMP and thus respects its settings.)
Alternatively, this module provides :func:`omp_get_nthreads` and :func:`omp_set_nthreads`,
which give their best to get and set the number of threads at runtime,
while still being failsafe if the shared OpenMP library is not found. In the latter case,
you might also try the equivalent :func:`mkl_get_nthreads` and :func:`mkl_set_nthreads`.
"""
# Copyright (C) TeNPy Developers, GNU GPLv3
import warnings
import ctypes
from ctypes.util import find_library
import sys
__all__ = [
'memory_usage', 'load_omp_library', 'omp_get_nthreads', 'omp_set_nthreads', 'mkl_get_nthreads',
'mkl_set_nthreads'
]
_omp_lib = None
def memory_usage():
"""Return memory usage of the running python process.
You can ``pip install psutil`` if you get only ``-1.``.
Returns
-------
mem : float
Currently used memory in megabytes. ``-1.`` if no way to read out.
"""
try:
import resource # linux-only
unit_to_MB = 1024**2 if sys.platform == 'darwin' else 1024 # linux uses kB, but MacOS byte
# see also issue #262
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / unit_to_MB
except ImportError:
pass
try:
import psutil
proc = psutil.Process()
return proc.memory_info().rss / 1024**2
except ImportError:
pass
warnings.warn("No tool to determine memory_usage")
return -1.
def load_omp_library(libs=["libiomp5.so",
find_library("libiomp5md"),
find_library("gomp")],
verbose=True):
"""Tries to load openMP library.
Parameters
----------
libs :
list of possible library names we should try to load (with ctypes.CDLL).
verbose : bool
whether to print the name of the loaded library.
Returns
-------
omp : CDLL | None
OpenMP shared library if found, otherwise None.
Once it was successfully imported, no re-imports are tried.
"""
global _omp_lib
if _omp_lib is None:
for l in libs:
if l is None:
continue
try:
_omp_lib = ctypes.CDLL(l)
if verbose:
print("loaded " + l + " for omp")
break
except OSError:
pass
if _omp_lib is None:
warnings.warn("OpenMP library not found: can't get/set nthreads")
return _omp_lib
def omp_get_nthreads():
"""wrapper around OpenMP ``get_max_threads``.
Returns
-------
max_threads : int
The maximum number of threads used by OpenMP (and thus MKL).
``-1`` if unable to read out.
"""
omp = load_omp_library()
if omp is not None:
return omp.omp_get_max_threads()
return -1
def omp_set_nthreads(n):
"""wrapper around OpenMP ``set_nthreads``.
Parameters
----------
n : int
the number of threads to use
Returns
-------
success : bool
whether the shared library was found and set.
"""
omp = load_omp_library()
if omp is not None:
omp.omp_set_num_threads(int(n))
return True
return False
def mkl_get_nthreads():
"""wrapper around MKL ``get_max_threads``.
Returns
-------
max_threads : int
The maximum number of threads used by MKL. ``-1`` if unable to read out.
"""
try:
import mkl # available in conda MKL
return mkl.get_max_threads()
except ImportError:
try:
mkl_rt = ctypes.CDLL('libmkl_rt.so')
return mkl_rt.mkl_get_max_threads()
except OSError:
warnings.warn("MKL library not found: can't get nthreads")
return -1
def mkl_set_nthreads(n):
"""wrapper around MKL ``set_num_threads``.
Parameters
----------
n : int
the number of threads to use
Returns
-------
success : bool
whether the shared library was found and set.
"""
try:
import mkl # available in conda MKL
mkl.set_num_threads(n)
return True
except ImportError:
try:
mkl_rt = ctypes.CDLL('libmkl_rt.so')
mkl_rt.mkl_set_num_threads(ctypes.byref(ctypes.c_int(n)))
return True
except OSError:
warnings.warn("MKL library not found: can't set nthreads")
return False