-
Notifications
You must be signed in to change notification settings - Fork 45
/
context.py
378 lines (301 loc) · 13.3 KB
/
context.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
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
###############################################################################
# Basic context management with LokyContext
#
# author: Thomas Moreau and Olivier Grisel
#
# adapted from multiprocessing/context.py
# * Create a context ensuring loky uses only objects that are compatible
# * Add LokyContext to the list of context of multiprocessing so loky can be
# used with multiprocessing.set_start_method
# * Implement a CFS-aware amd physical-core aware cpu_count function.
#
import os
import sys
import math
import subprocess
import traceback
import warnings
import multiprocessing as mp
from multiprocessing import get_context as mp_get_context
from multiprocessing.context import BaseContext
from .process import LokyProcess, LokyInitMainProcess
# Apparently, on older Python versions, loky cannot work 61 workers on Windows
# but instead 60: ¯\_(ツ)_/¯
if sys.version_info >= (3, 8):
from concurrent.futures.process import _MAX_WINDOWS_WORKERS
if sys.version_info < (3, 10):
_MAX_WINDOWS_WORKERS = _MAX_WINDOWS_WORKERS - 1
else:
# compat for versions before 3.8 which do not define this.
_MAX_WINDOWS_WORKERS = 60
START_METHODS = ["loky", "loky_init_main", "spawn"]
if sys.platform != "win32":
START_METHODS += ["fork", "forkserver"]
_DEFAULT_START_METHOD = None
# Cache for the number of physical cores to avoid repeating subprocess calls.
# It should not change during the lifetime of the program.
physical_cores_cache = None
def get_context(method=None):
# Try to overload the default context
method = method or _DEFAULT_START_METHOD or "loky"
if method == "fork":
# If 'fork' is explicitly requested, warn user about potential issues.
warnings.warn(
"`fork` start method should not be used with "
"`loky` as it does not respect POSIX. Try using "
"`spawn` or `loky` instead.",
UserWarning,
)
try:
return mp_get_context(method)
except ValueError:
raise ValueError(
f"Unknown context '{method}'. Value should be in "
f"{START_METHODS}."
)
def set_start_method(method, force=False):
global _DEFAULT_START_METHOD
if _DEFAULT_START_METHOD is not None and not force:
raise RuntimeError("context has already been set")
assert method is None or method in START_METHODS, (
f"'{method}' is not a valid start_method. It should be in "
f"{START_METHODS}"
)
_DEFAULT_START_METHOD = method
def get_start_method():
return _DEFAULT_START_METHOD
def cpu_count(only_physical_cores=False):
"""Return the number of CPUs the current process can use.
The returned number of CPUs accounts for:
* the number of CPUs in the system, as given by
``multiprocessing.cpu_count``;
* the CPU affinity settings of the current process
(available on some Unix systems);
* Cgroup CPU bandwidth limit (available on Linux only, typically
set by docker and similar container orchestration systems);
* the value of the LOKY_MAX_CPU_COUNT environment variable if defined.
and is given as the minimum of these constraints.
If ``only_physical_cores`` is True, return the number of physical cores
instead of the number of logical cores (hyperthreading / SMT). Note that
this option is not enforced if the number of usable cores is controlled in
any other way such as: process affinity, Cgroup restricted CPU bandwidth
or the LOKY_MAX_CPU_COUNT environment variable. If the number of physical
cores is not found, return the number of logical cores.
Note that on Windows, the returned number of CPUs cannot exceed 61 (or 60 for
Python < 3.10), see:
https://bugs.python.org/issue26903.
It is also always larger or equal to 1.
"""
# Note: os.cpu_count() is allowed to return None in its docstring
os_cpu_count = os.cpu_count() or 1
if sys.platform == "win32":
# On Windows, attempting to use more than 61 CPUs would result in a
# OS-level error. See https://bugs.python.org/issue26903. According to
# https://learn.microsoft.com/en-us/windows/win32/procthread/processor-groups
# it might be possible to go beyond with a lot of extra work but this
# does not look easy.
os_cpu_count = min(os_cpu_count, _MAX_WINDOWS_WORKERS)
cpu_count_user = _cpu_count_user(os_cpu_count)
aggregate_cpu_count = max(min(os_cpu_count, cpu_count_user), 1)
if not only_physical_cores:
return aggregate_cpu_count
if cpu_count_user < os_cpu_count:
# Respect user setting
return max(cpu_count_user, 1)
cpu_count_physical, exception = _count_physical_cores()
if cpu_count_physical != "not found":
return cpu_count_physical
# Fallback to default behavior
if exception is not None:
# warns only the first time
warnings.warn(
"Could not find the number of physical cores for the "
f"following reason:\n{exception}\n"
"Returning the number of logical cores instead. You can "
"silence this warning by setting LOKY_MAX_CPU_COUNT to "
"the number of cores you want to use."
)
traceback.print_tb(exception.__traceback__)
return aggregate_cpu_count
def _cpu_count_cgroup(os_cpu_count):
# Cgroup CPU bandwidth limit available in Linux since 2.6 kernel
cpu_max_fname = "/sys/fs/cgroup/cpu.max"
cfs_quota_fname = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"
cfs_period_fname = "/sys/fs/cgroup/cpu/cpu.cfs_period_us"
if os.path.exists(cpu_max_fname):
# cgroup v2
# https://www.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html
with open(cpu_max_fname) as fh:
cpu_quota_us, cpu_period_us = fh.read().strip().split()
elif os.path.exists(cfs_quota_fname) and os.path.exists(cfs_period_fname):
# cgroup v1
# https://www.kernel.org/doc/html/latest/scheduler/sched-bwc.html#management
with open(cfs_quota_fname) as fh:
cpu_quota_us = fh.read().strip()
with open(cfs_period_fname) as fh:
cpu_period_us = fh.read().strip()
else:
# No Cgroup CPU bandwidth limit (e.g. non-Linux platform)
cpu_quota_us = "max"
cpu_period_us = 100_000 # unused, for consistency with default values
if cpu_quota_us == "max":
# No active Cgroup quota on a Cgroup-capable platform
return os_cpu_count
else:
cpu_quota_us = int(cpu_quota_us)
cpu_period_us = int(cpu_period_us)
if cpu_quota_us > 0 and cpu_period_us > 0:
return math.ceil(cpu_quota_us / cpu_period_us)
else: # pragma: no cover
# Setting a negative cpu_quota_us value is a valid way to disable
# cgroup CPU bandwith limits
return os_cpu_count
def _cpu_count_affinity(os_cpu_count):
# Number of available CPUs given affinity settings
if hasattr(os, "sched_getaffinity"):
try:
return len(os.sched_getaffinity(0))
except NotImplementedError:
pass
# On PyPy and possibly other platforms, os.sched_getaffinity does not exist
# or raises NotImplementedError, let's try with the psutil if installed.
try:
import psutil
p = psutil.Process()
if hasattr(p, "cpu_affinity"):
return len(p.cpu_affinity())
except ImportError: # pragma: no cover
if (
sys.platform == "linux"
and os.environ.get("LOKY_MAX_CPU_COUNT") is None
):
# PyPy does not implement os.sched_getaffinity on Linux which
# can cause severe oversubscription problems. Better warn the
# user in this particularly pathological case which can wreck
# havoc, typically on CI workers.
warnings.warn(
"Failed to inspect CPU affinity constraints on this system. "
"Please install psutil or explictly set LOKY_MAX_CPU_COUNT."
)
# This can happen for platforms that do not implement any kind of CPU
# infinity such as macOS-based platforms.
return os_cpu_count
def _cpu_count_user(os_cpu_count):
"""Number of user defined available CPUs"""
cpu_count_affinity = _cpu_count_affinity(os_cpu_count)
cpu_count_cgroup = _cpu_count_cgroup(os_cpu_count)
# User defined soft-limit passed as a loky specific environment variable.
cpu_count_loky = int(os.environ.get("LOKY_MAX_CPU_COUNT", os_cpu_count))
return min(cpu_count_affinity, cpu_count_cgroup, cpu_count_loky)
def _count_physical_cores():
"""Return a tuple (number of physical cores, exception)
If the number of physical cores is found, exception is set to None.
If it has not been found, return ("not found", exception).
The number of physical cores is cached to avoid repeating subprocess calls.
"""
exception = None
# First check if the value is cached
global physical_cores_cache
if physical_cores_cache is not None:
return physical_cores_cache, exception
# Not cached yet, find it
try:
if sys.platform == "linux":
cpu_info = subprocess.run(
"lscpu --parse=core".split(), capture_output=True, text=True
)
cpu_info = cpu_info.stdout.splitlines()
cpu_info = {line for line in cpu_info if not line.startswith("#")}
cpu_count_physical = len(cpu_info)
elif sys.platform == "win32":
cpu_info = subprocess.run(
"wmic CPU Get NumberOfCores /Format:csv".split(),
capture_output=True,
text=True,
)
cpu_info = cpu_info.stdout.splitlines()
cpu_info = [
l.split(",")[1]
for l in cpu_info
if (l and l != "Node,NumberOfCores")
]
cpu_count_physical = sum(map(int, cpu_info))
elif sys.platform == "darwin":
cpu_info = subprocess.run(
"sysctl -n hw.physicalcpu".split(),
capture_output=True,
text=True,
)
cpu_info = cpu_info.stdout
cpu_count_physical = int(cpu_info)
else:
raise NotImplementedError(f"unsupported platform: {sys.platform}")
# if cpu_count_physical < 1, we did not find a valid value
if cpu_count_physical < 1:
raise ValueError(f"found {cpu_count_physical} physical cores < 1")
except Exception as e:
exception = e
cpu_count_physical = "not found"
# Put the result in cache
physical_cores_cache = cpu_count_physical
return cpu_count_physical, exception
class LokyContext(BaseContext):
"""Context relying on the LokyProcess."""
_name = "loky"
Process = LokyProcess
cpu_count = staticmethod(cpu_count)
def Queue(self, maxsize=0, reducers=None):
"""Returns a queue object"""
from .queues import Queue
return Queue(maxsize, reducers=reducers, ctx=self.get_context())
def SimpleQueue(self, reducers=None):
"""Returns a queue object"""
from .queues import SimpleQueue
return SimpleQueue(reducers=reducers, ctx=self.get_context())
if sys.platform != "win32":
"""For Unix platform, use our custom implementation of synchronize
ensuring that we use the loky.backend.resource_tracker to clean-up
the semaphores in case of a worker crash.
"""
def Semaphore(self, value=1):
"""Returns a semaphore object"""
from .synchronize import Semaphore
return Semaphore(value=value)
def BoundedSemaphore(self, value):
"""Returns a bounded semaphore object"""
from .synchronize import BoundedSemaphore
return BoundedSemaphore(value)
def Lock(self):
"""Returns a lock object"""
from .synchronize import Lock
return Lock()
def RLock(self):
"""Returns a recurrent lock object"""
from .synchronize import RLock
return RLock()
def Condition(self, lock=None):
"""Returns a condition object"""
from .synchronize import Condition
return Condition(lock)
def Event(self):
"""Returns an event object"""
from .synchronize import Event
return Event()
class LokyInitMainContext(LokyContext):
"""Extra context with LokyProcess, which does load the main module
This context is used for compatibility in the case ``cloudpickle`` is not
present on the running system. This permits to load functions defined in
the ``main`` module, using proper safeguards. The declaration of the
``executor`` should be protected by ``if __name__ == "__main__":`` and the
functions and variable used from main should be out of this block.
This mimics the default behavior of multiprocessing under Windows and the
behavior of the ``spawn`` start method on a posix system.
For more details, see the end of the following section of python doc
https://docs.python.org/3/library/multiprocessing.html#multiprocessing-programming
"""
_name = "loky_init_main"
Process = LokyInitMainProcess
# Register loky context so it works with multiprocessing.get_context
ctx_loky = LokyContext()
mp.context._concrete_contexts["loky"] = ctx_loky
mp.context._concrete_contexts["loky_init_main"] = LokyInitMainContext()