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 

11 

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

13 

14 

15def _try_extract_cgroup_cpu_quota(): 

16 # cgroup v1 

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

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

19 try: 

20 with open("/sys/fs/cgroup/%s/cpu.cfs_quota_us" % dirname) as f: 

21 quota = int(f.read()) 

22 with open("/sys/fs/cgroup/%s/cpu.cfs_period_us" % dirname) as f: 

23 period = int(f.read()) 

24 return quota, period 

25 except Exception: 

26 pass 

27 

28 # cgroup v2 

29 try: 

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

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

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

33 group_path = f"{group_path}/" 

34 with open("/sys/fs/cgroup%scpu.max" % group_path) as f: 

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

36 return quota, period 

37 except Exception: 

38 pass 

39 

40 # No cgroup CPU quota found 

41 return None, None 

42 

43 

44def cpu_count(): 

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

46 

47 Takes the minimum value from the following locations: 

48 

49 - Total system cpus available on the host. 

50 - CPU Affinity (if set) 

51 - Cgroups limit (if set) 

52 """ 

53 count = os.cpu_count() 

54 

55 # Check CPU affinity if available 

56 if psutil is not None: 

57 try: 

58 affinity_count = len(psutil.Process().cpu_affinity()) 

59 if affinity_count > 0: 

60 count = min(count, affinity_count) 

61 except Exception: 

62 pass 

63 

64 # Check cgroups if available 

65 if sys.platform == "linux": 

66 quota, period = _try_extract_cgroup_cpu_quota() 

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

68 # We round up on fractional CPUs 

69 cgroups_count = math.ceil(quota / period) 

70 if cgroups_count > 0: 

71 count = min(count, cgroups_count) 

72 

73 return count 

74 

75 

76CPU_COUNT = cpu_count()