Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1###############################################################################
2# Basic context management with LokyContext
3#
4# author: Thomas Moreau and Olivier Grisel
5#
6# adapted from multiprocessing/context.py
7# * Create a context ensuring loky uses only objects that are compatible
8# * Add LokyContext to the list of context of multiprocessing so loky can be
9# used with multiprocessing.set_start_method
10# * Implement a CFS-aware amd physical-core aware cpu_count function.
11#
12import ctypes
13import math
14import multiprocessing as mp
15import os
16import subprocess
17import sys
18import traceback
19import warnings
21from ctypes import wintypes
22from multiprocessing import get_context as mp_get_context
23from multiprocessing.context import BaseContext
24from concurrent.futures.process import _MAX_WINDOWS_WORKERS
27from .process import LokyProcess, LokyInitMainProcess
29# Apparently, on older Python versions, loky cannot work 61 workers on Windows
30# but instead 60: ¯\_(ツ)_/¯
31if sys.version_info < (3, 10):
32 _MAX_WINDOWS_WORKERS = _MAX_WINDOWS_WORKERS - 1
34START_METHODS = ["loky", "loky_init_main", "spawn"]
35if sys.platform != "win32":
36 START_METHODS += ["fork", "forkserver"]
38_DEFAULT_START_METHOD = None
40# Cache for the number of physical cores, to avoid repeating subprocess
41# calls or re-parsing /proc/cpuinfo. Keyed by None for the whole machine, or
42# by a frozenset of logical CPU ids when restricted to a CPU affinity mask,
43# see `_count_physical_cores`.
44physical_cores_cache = {}
47def get_context(method=None):
48 # Try to overload the default context
49 method = method or _DEFAULT_START_METHOD or "loky"
50 if method == "fork":
51 # If 'fork' is explicitly requested, warn user about potential issues.
52 warnings.warn(
53 "`fork` start method should not be used with "
54 "`loky` as it does not respect POSIX. Try using "
55 "`spawn` or `loky` instead.",
56 UserWarning,
57 )
58 try:
59 return mp_get_context(method)
60 except ValueError:
61 raise ValueError(
62 f"Unknown context '{method}'. Value should be in "
63 f"{START_METHODS}."
64 )
67def set_start_method(method, force=False):
68 global _DEFAULT_START_METHOD
69 if _DEFAULT_START_METHOD is not None and not force:
70 raise RuntimeError("context has already been set")
71 assert method is None or method in START_METHODS, (
72 f"'{method}' is not a valid start_method. It should be in "
73 f"{START_METHODS}"
74 )
76 _DEFAULT_START_METHOD = method
79def get_start_method():
80 return _DEFAULT_START_METHOD
83def cpu_count(only_physical_cores=False):
84 """Return the number of CPUs the current process can use.
86 The returned number of CPUs accounts for:
87 * the number of CPUs in the system, as given by
88 ``multiprocessing.cpu_count``;
89 * the CPU affinity settings of the current process
90 (available on some Unix systems);
91 * Cgroup CPU bandwidth limit (available on Linux only, typically
92 set by docker and similar container orchestration systems);
93 * the value of the LOKY_MAX_CPU_COUNT environment variable if defined.
94 and is given as the minimum of these constraints.
96 If ``only_physical_cores`` is True, return the number of physical cores
97 instead of the number of logical cores (hyperthreading / SMT): this is
98 given as the minimum of the same constraints as above, together with the
99 number of physical cores.
100 On Linux, the computed number of physical cores is restrained to the
101 physical cores reachable through the CPU affinity mask.
102 On other platforms, it's the number of physical cores on the whole
103 machine, since the affinity mask cannot be taken into account there.
105 Note that on Windows, the returned number of CPUs cannot exceed 61, see:
106 https://bugs.python.org/issue26903.
108 It is also always larger or equal to 1.
109 """
110 # Note: os.cpu_count() is allowed to return None in its docstring
111 os_cpu_count = os.cpu_count() or 1
112 if sys.platform == "win32":
113 # On Windows, attempting to use more than 61 CPUs would result in a
114 # OS-level error. See https://bugs.python.org/issue26903. According to
115 # https://learn.microsoft.com/en-us/windows/win32/procthread/processor-groups
116 # it might be possible to go beyond with a lot of extra work but this
117 # does not look easy.
118 os_cpu_count = min(os_cpu_count, _MAX_WINDOWS_WORKERS)
120 cpu_affinity_set = _cpu_count_affinity_set()
121 cpu_count_affinity = (
122 os_cpu_count if cpu_affinity_set is None else len(cpu_affinity_set)
123 )
124 cpu_count_cgroup = _cpu_count_cgroup(os_cpu_count)
125 cpu_count_loky = int(os.environ.get("LOKY_MAX_CPU_COUNT", os_cpu_count))
127 if not only_physical_cores:
128 return max(
129 1,
130 min(
131 os_cpu_count,
132 cpu_count_affinity,
133 cpu_count_cgroup,
134 cpu_count_loky,
135 ),
136 )
138 cpu_count_physical, exception = _count_physical_cores(
139 # On Linux, try to collapse SMT/hyper-threading sibling logical CPUs
140 # reachable through the current CPU affinity mask that share the same
141 # physical core. See https://github.com/joblib/loky/issues/639.
142 # On other platforms we lack easy access to CPU topology info to refine
143 # the count, so just pass None for `cpu_affinity_set`.
144 cpu_affinity_set
145 if sys.platform == "linux"
146 else None
147 )
148 if cpu_count_physical == "not found":
149 cpu_count_physical = os_cpu_count
151 _warn_physical_cores_not_found(exception)
153 return max(
154 1,
155 min(
156 os_cpu_count,
157 cpu_count_physical,
158 cpu_count_affinity,
159 cpu_count_cgroup,
160 cpu_count_loky,
161 ),
162 )
165def _warn_physical_cores_not_found(exception):
166 if exception is not None:
167 # warns only the first time
168 warnings.warn(
169 "Could not find the number of physical cores for the "
170 f"following reason:\n{exception}\n"
171 "Returning the number of logical cores instead. You can "
172 "silence this warning by setting LOKY_MAX_CPU_COUNT to "
173 "the number of cores you want to use."
174 )
175 traceback.print_tb(exception.__traceback__)
178def _cpu_count_cgroup(os_cpu_count):
179 # Cgroup CPU bandwidth limit available in Linux since 2.6 kernel
180 cpu_max_fname = "/sys/fs/cgroup/cpu.max"
181 cfs_quota_fname = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"
182 cfs_period_fname = "/sys/fs/cgroup/cpu/cpu.cfs_period_us"
184 cpu_quota_us = None
185 cpu_period_us = None
187 if os.path.exists(cpu_max_fname):
188 # cgroup v2
189 # https://www.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html
190 with open(cpu_max_fname) as fh:
191 # Parse the quota and period values
192 parts = fh.read().strip().split()
193 if len(parts) == 2:
194 cpu_quota_us, cpu_period_us = parts
195 # If len(parts) != 2, leave as None and fall back to v1
197 # If we didn't get values from cgroup v2, try cgroup v1
198 if cpu_quota_us is None or cpu_period_us is None:
199 if os.path.exists(cfs_quota_fname) and os.path.exists(
200 cfs_period_fname
201 ):
202 # cgroup v1
203 # https://www.kernel.org/doc/html/latest/scheduler/sched-bwc.html#management
204 with open(cfs_quota_fname) as fh:
205 cpu_quota_us = fh.read().strip()
206 with open(cfs_period_fname) as fh:
207 cpu_period_us = fh.read().strip()
208 else:
209 # No Cgroup CPU bandwidth limit (e.g. non-Linux platform)
210 cpu_quota_us = "max"
212 if cpu_quota_us == "max":
213 # No active Cgroup quota on a Cgroup-capable platform
214 return os_cpu_count
215 else:
216 cpu_quota_us = int(cpu_quota_us)
217 cpu_period_us = int(cpu_period_us)
218 if cpu_quota_us > 0 and cpu_period_us > 0:
219 return math.ceil(cpu_quota_us / cpu_period_us)
220 else: # pragma: no cover
221 # Setting a negative cpu_quota_us value is a valid way to disable
222 # cgroup CPU bandwidth limits
223 return os_cpu_count
226def _cpu_count_affinity_set():
227 """Return the current CPU affinity mask as a set of logical CPU ids.
229 Return None if the affinity mask cannot be determined on this platform,
230 for instance because neither `os.sched_getaffinity` nor `psutil` are
231 available.
232 """
233 if hasattr(os, "sched_getaffinity"):
234 try:
235 return os.sched_getaffinity(0)
236 except NotImplementedError:
237 pass
239 # On some platforms, os.sched_getaffinity does not exist or raises
240 # NotImplementedError, let's try with the psutil if installed.
241 try:
242 import psutil
244 p = psutil.Process()
245 if hasattr(p, "cpu_affinity"):
246 return set(p.cpu_affinity())
248 except ImportError: # pragma: no cover
249 if (
250 sys.platform == "linux"
251 and os.environ.get("LOKY_MAX_CPU_COUNT") is None
252 ):
253 # Some platforms don't implement os.sched_getaffinity on Linux which
254 # can cause severe oversubscription problems. Better warn the
255 # user in this particularly pathological case which can wreck
256 # havoc, typically on CI workers.
257 warnings.warn(
258 "Failed to inspect CPU affinity constraints on this system. "
259 "Please install psutil or explictly set LOKY_MAX_CPU_COUNT."
260 )
262 return None
265def _count_physical_cores(cpu_set=None):
266 """Return a tuple (number of physical cores, exception)
268 If the number of physical cores is found, exception is set to None.
269 If it has not been found, return ("not found", exception).
271 If `cpu_set` is not None, only the logical CPUs it contains are
272 considered when counting physical cores, which also collapses
273 SMT/hyper-threading sibling logical CPUs that share the same physical
274 core. Only implemented for Linux, where per-CPU topology information is
275 readily available; callers are expected to only pass a non-None
276 `cpu_set` on Linux.
278 The number of physical cores is cached (per distinct `cpu_set` when one
279 is passed) to avoid repeating subprocess calls or re-parsing
280 /proc/cpuinfo.
281 """
282 exception = None
284 # First check if the value is cached
285 global physical_cores_cache
286 cache_key = None if cpu_set is None else frozenset(cpu_set)
287 if cache_key in physical_cores_cache:
288 return physical_cores_cache[cache_key], None
290 # Not cached yet, find it
291 try:
292 if sys.platform == "linux":
293 cpu_count_physical = _count_physical_cores_linux(cpu_set)
294 elif cpu_set is not None:
295 raise NotImplementedError(
296 "affinity-aware physical core counting is only "
297 "implemented on Linux"
298 )
299 elif sys.platform == "win32":
300 cpu_count_physical = _count_physical_cores_win32()
301 elif sys.platform == "darwin":
302 cpu_count_physical = _count_physical_cores_darwin()
303 elif sys.platform.startswith("freebsd"):
304 cpu_count_physical = _count_physical_cores_freebsd()
305 else:
306 raise NotImplementedError(f"unsupported platform: {sys.platform}")
308 # if cpu_count_physical < 1, we did not find a valid value
309 if cpu_count_physical < 1:
310 raise ValueError(f"found {cpu_count_physical} physical cores < 1")
312 except Exception as e:
313 exception = e
314 cpu_count_physical = "not found"
316 # Put the result in cache
317 physical_cores_cache[cache_key] = cpu_count_physical
319 return cpu_count_physical, exception
322def _count_physical_cores_linux(cpu_set=None):
323 """Return the number of distinct physical cores on Linux, by parsing
324 /proc/cpuinfo.
326 A physical core is identified by its (physical id, core id) pair: core
327 id alone is only guaranteed unique within a physical package, so a
328 multi-socket machine can otherwise under-count cores that share the
329 same core id across sockets.
331 If `cpu_set` is not None, only the logical CPUs it contains are
332 considered, which also collapses SMT/hyper-threading sibling logical
333 CPUs that share the same physical core.
334 """
335 with open("/proc/cpuinfo") as f:
336 cpuinfo = f.read()
338 cores = set()
339 processor = core_id = physical_id = None
340 for line in cpuinfo.splitlines() + [""]:
341 if not line.strip():
342 # blank line: end of the current logical CPU block
343 if core_id is not None and (
344 cpu_set is None or processor in cpu_set
345 ):
346 cores.add((physical_id, core_id))
347 processor = core_id = physical_id = None
348 continue
350 key, sep, value = line.partition(":")
351 if not sep:
352 continue
353 key = key.strip()
354 value = value.strip()
355 if key == "processor":
356 processor = int(value)
357 elif key == "core id":
358 core_id = value
359 elif key == "physical id":
360 physical_id = value
362 if not cores:
363 raise ValueError("could not find any physical core in /proc/cpuinfo")
365 return len(cores)
368def _count_physical_cores_win32():
369 try:
370 return _count_physical_cores_win32_ctypes()
371 except Exception:
372 pass # fallback to powershell
373 try:
374 return _count_physical_cores_win32_powershell()
375 except Exception:
376 pass # fallback to wmic (older Windows versions; deprecated now)
378 cpu_info = subprocess.run(
379 "wmic CPU Get NumberOfCores /Format:csv".split(),
380 capture_output=True,
381 text=True,
382 creationflags=subprocess.CREATE_NO_WINDOW,
383 )
384 cpu_info = cpu_info.stdout.splitlines()
385 cpu_info = [
386 l.split(",")[1] for l in cpu_info if (l and l != "Node,NumberOfCores")
387 ]
388 return sum(map(int, cpu_info))
391def _count_physical_cores_win32_powershell():
392 cmd = "-NoProfile -Command (Get-CimInstance -ClassName Win32_Processor).NumberOfCores"
393 cpu_info = subprocess.run(
394 f"powershell.exe {cmd}".split(),
395 capture_output=True,
396 text=True,
397 creationflags=subprocess.CREATE_NO_WINDOW,
398 )
399 cpu_info = cpu_info.stdout.splitlines()
400 return sum(map(int, cpu_info))
403def _count_physical_cores_win32_ctypes():
404 ERROR_INSUFFICIENT_BUFFER = 122
405 RelationProcessorCore = 0
407 kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
408 get_logical_processor_information = (
409 kernel32.GetLogicalProcessorInformationEx
410 )
411 get_logical_processor_information.argtypes = [
412 wintypes.DWORD,
413 ctypes.c_void_p,
414 ctypes.POINTER(wintypes.DWORD),
415 ]
416 get_logical_processor_information.restype = wintypes.BOOL
418 # Mirror the header of SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX. The full
419 # structure is variable-sized, and only Relationship and Size are needed.
420 # https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-system_logical_processor_information_ex
421 class SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX(ctypes.Structure):
422 _fields_ = [
423 ("Relationship", wintypes.DWORD),
424 ("Size", wintypes.DWORD),
425 ]
427 # First obtain the required buffer size. This call is expected to fail with
428 # ERROR_INSUFFICIENT_BUFFER and set returned_length.
429 # https://learn.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getlogicalprocessorinformationex
430 returned_length = wintypes.DWORD()
431 returned_length_ref = ctypes.byref(returned_length)
432 if get_logical_processor_information(
433 RelationProcessorCore, None, returned_length_ref
434 ):
435 raise RuntimeError("unexpected successful buffer sizing call")
437 error = ctypes.get_last_error()
438 if error != ERROR_INSUFFICIENT_BUFFER:
439 raise ctypes.WinError(error)
441 buf = ctypes.create_string_buffer(returned_length.value)
442 if not get_logical_processor_information(
443 RelationProcessorCore, buf, returned_length_ref
444 ):
445 raise ctypes.WinError(ctypes.get_last_error())
447 offset = 0
448 physical_core_count = 0
449 header_size = ctypes.sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)
451 while offset < returned_length.value:
452 remaining = returned_length.value - offset
453 if remaining < header_size:
454 raise RuntimeError("truncated processor information record")
456 processor_core_info = (
457 SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX.from_buffer(buf, offset)
458 )
459 record_size = processor_core_info.Size
460 if record_size < header_size or record_size > remaining:
461 raise RuntimeError("invalid processor information record size")
462 if processor_core_info.Relationship != RelationProcessorCore:
463 raise RuntimeError("unexpected logical processor relationship")
465 physical_core_count += 1
466 offset += record_size
468 if physical_core_count == 0:
469 raise RuntimeError("Windows reported no active physical cores")
471 if physical_core_count < 1:
472 raise RuntimeError(
473 "GetLogicalProcessorInformationEx returned no physical cores"
474 )
475 return physical_core_count
478def _count_physical_cores_darwin():
479 cpu_info = subprocess.run(
480 "sysctl -n hw.physicalcpu".split(),
481 capture_output=True,
482 text=True,
483 )
484 cpu_info = cpu_info.stdout
485 return int(cpu_info)
488def _count_physical_cores_freebsd():
489 cpu_info = subprocess.run(
490 "sysctl -n kern.smp.cores".split(),
491 capture_output=True,
492 text=True,
493 )
494 cpu_info = cpu_info.stdout
495 return int(cpu_info)
498class LokyContext(BaseContext):
499 """Context relying on the LokyProcess."""
501 _name = "loky"
502 Process = LokyProcess
503 cpu_count = staticmethod(cpu_count)
505 def Queue(self, maxsize=0, reducers=None):
506 """Returns a queue object"""
507 from .queues import Queue
509 return Queue(maxsize, reducers=reducers, ctx=self.get_context())
511 def SimpleQueue(self, reducers=None):
512 """Returns a queue object"""
513 from .queues import SimpleQueue
515 return SimpleQueue(reducers=reducers, ctx=self.get_context())
517 if sys.platform != "win32":
518 """For Unix platform, use our custom implementation of synchronize
519 ensuring that we use the loky.backend.resource_tracker to clean-up
520 the semaphores in case of a worker crash.
521 """
523 def Semaphore(self, value=1):
524 """Returns a semaphore object"""
525 from .synchronize import Semaphore
527 return Semaphore(value=value)
529 def BoundedSemaphore(self, value):
530 """Returns a bounded semaphore object"""
531 from .synchronize import BoundedSemaphore
533 return BoundedSemaphore(value)
535 def Lock(self):
536 """Returns a lock object"""
537 from .synchronize import Lock
539 return Lock()
541 def RLock(self):
542 """Returns a recurrent lock object"""
543 from .synchronize import RLock
545 return RLock()
547 def Condition(self, lock=None):
548 """Returns a condition object"""
549 from .synchronize import Condition
551 return Condition(lock)
553 def Event(self):
554 """Returns an event object"""
555 from .synchronize import Event
557 return Event()
560class LokyInitMainContext(LokyContext):
561 """Extra context with LokyProcess, which does load the main module
563 This context is used for compatibility in the case ``cloudpickle`` is not
564 present on the running system. This permits to load functions defined in
565 the ``main`` module, using proper safeguards. The declaration of the
566 ``executor`` should be protected by ``if __name__ == "__main__":`` and the
567 functions and variable used from main should be out of this block.
569 This mimics the default behavior of multiprocessing under Windows and the
570 behavior of the ``spawn`` start method on a posix system.
571 For more details, see the end of the following section of python doc
572 https://docs.python.org/3/library/multiprocessing.html#multiprocessing-programming
573 """
575 _name = "loky_init_main"
576 Process = LokyInitMainProcess
579# Register loky context so it works with multiprocessing.get_context
580ctx_loky = LokyContext()
581mp.context._concrete_contexts["loky"] = ctx_loky
582mp.context._concrete_contexts["loky_init_main"] = LokyInitMainContext()