Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/dask/system.py: 62%

1from __future__ import annotations 

2 

3import math 

4import os 

5import sys 

6 

7try: 

8 import psutil 

9except ImportError: 

10 psutil = None # type: ignore[assignment] 

11 

12from dask._compatibility import LINUX 

13 

14__all__ = ("cpu_count", "CPU_COUNT") 

15 

16 

17def _try_extract_cgroup_cpu_quota(): 

18 # cgroup v1 

19 # The directory name isn't standardized across linux distros, check both 

20 for dirname in ["cpuacct,cpu", "cpu,cpuacct"]: 

21 try: 

22 with open(f"/sys/fs/cgroup/{dirname}/cpu.cfs_quota_us") as f: 

23 quota = int(f.read()) 

24 with open(f"/sys/fs/cgroup/{dirname}/cpu.cfs_period_us") as f: 

25 period = int(f.read()) 

26 return quota, period 

27 except Exception: 

28 pass 

29 

30 # cgroup v2 

31 try: 

32 with open("/proc/self/cgroup") as f: 

33 group_path = f.read().strip().split(":")[-1] 

34 if not group_path.endswith("/"): 

35 group_path = f"{group_path}/" 

36 with open(f"/sys/fs/cgroup{group_path}cpu.max") as f: 

37 quota, period = map(int, f.read().split(" ")) 

38 return quota, period 

39 except Exception: 

40 pass 

41 

42 # No cgroup CPU quota found 

43 return None, None 

44 

45 

46def cpu_count() -> int: 

47 """Get the available CPU count for this system. 

48 

49 Takes the minimum value from the following locations: 

50 

51 - Total system cpus available on the host. 

52 - CPU Affinity (if set) 

53 - Cgroups limit (if set) 

54 """ 

55 if sys.version_info >= (3, 13): 

56 # Embeds CPU affinity checks 

57 count = os.process_cpu_count() 

58 elif hasattr(os, "sched_getaffinity"): 

59 # https://docs.python.org/3/library/os.html#interface-to-the-scheduler 

60 # "only available on some Unix platforms"; neither MacOS nor Windows 

61 count = len(os.sched_getaffinity(0)) 

62 else: 

63 # Does not account for CPU affinity. 

64 # On exotic alternative Python implementations, it may return None. 

65 count = os.cpu_count() or 1 

66 assert count 

67 

68 # Additional CPU affinity check with psutil. 

69 # NOTE: do not limit this to Python <3.13: on Windows, 

70 # `psutil.Process().cpu_affinity(value)` does not change the reading of 

71 # os.process_cpu_count(). 

72 if psutil is not None: 

73 proc = psutil.Process() 

74 if hasattr(proc, "cpu_affinity"): 

75 affinity = proc.cpu_affinity() 

76 if affinity is not None: 

77 assert affinity 

78 count = min(count, len(affinity)) 

79 

80 # Check cgroups if available 

81 if LINUX: 

82 quota, period = _try_extract_cgroup_cpu_quota() 

83 if quota is not None and period is not None: 

84 # We round up on fractional CPUs 

85 cgroups_count = math.ceil(quota / period) 

86 if cgroups_count > 0: 

87 count = min(count, cgroups_count) 

88 

89 return count 

90 

91 

92CPU_COUNT = cpu_count()