Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/compat/_cpu.py: 20%

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

49 statements  

1""" 

2Detection of the CPU allocation actually available to this process. 

3 

4Used to bound the default worker count for parallel I/O, so that a default 

5:func:`~pandas.read_csv` does not oversubscribe a machine on which the process 

6has been given only a slice of the CPUs -- a CPU affinity mask, or a cgroup CPU 

7quota as set by ``docker --cpus`` and Kubernetes CPU limits. 

8""" 

9 

10from __future__ import annotations 

11 

12import functools 

13import os 

14 

15 

16def _read_sysfs_int(path: str) -> int | None: 

17 try: 

18 with open(path, encoding="utf-8") as fh: 

19 return int(fh.read().strip()) 

20 except (OSError, ValueError): 

21 return None 

22 

23 

24def _read_sysfs_str(path: str) -> str | None: 

25 # ValueError as well as OSError, matching _read_sysfs_int: a non-UTF-8 byte 

26 # raises UnicodeDecodeError, and this feeds available_cpu_count(), which 

27 # _default_n_workers calls unguarded -- so it would surface out of read_csv 

28 # rather than falling back. 

29 try: 

30 with open(path, encoding="utf-8") as fh: 

31 return fh.read().strip() 

32 except (OSError, ValueError): 

33 return None 

34 

35 

36def _parse_cgroup_v2_quota(text: str) -> float | None: 

37 """ 

38 CPU count from a cgroup v2 ``cpu.max`` value. 

39 

40 The format is ``"<quota> <period>"`` (both in microseconds) or ``"max"`` when 

41 the CPU time is unlimited; returns ``None`` when unlimited or malformed. 

42 """ 

43 parts = text.split() 

44 if not parts or parts[0] == "max": 

45 return None 

46 try: 

47 quota = int(parts[0]) 

48 period = int(parts[1]) if len(parts) > 1 else 100000 

49 except ValueError: 

50 return None 

51 if quota > 0 and period > 0: 

52 return quota / period 

53 return None 

54 

55 

56def _cgroup_cpu_quota() -> float | None: 

57 """CPUs allowed by the process's cgroup CFS quota, or None if unlimited.""" 

58 # cgroup v2; under a cgroup namespace (the common container case) this is the 

59 # limit the container itself sees. 

60 text = _read_sysfs_str("/sys/fs/cgroup/cpu.max") 

61 if text is not None: 

62 return _parse_cgroup_v2_quota(text) 

63 # cgroup v1 

64 quota = _read_sysfs_int("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") 

65 period = _read_sysfs_int("/sys/fs/cgroup/cpu/cpu.cfs_period_us") 

66 if quota and quota > 0 and period and period > 0: 

67 return quota / period 

68 return None 

69 

70 

71@functools.lru_cache(maxsize=1) 

72def available_cpu_count() -> int | None: 

73 """ 

74 Number of CPUs the process may actually use, or None if unconstrained. 

75 

76 Takes the tighter of the CPU affinity mask (``taskset``, SLURM cpusets) and 

77 the cgroup CFS quota (Docker ``--cpus``, Kubernetes CPU limits), so a default 

78 worker count does not oversubscribe a restricted environment. Returns 

79 ``None`` only where neither limit is probed -- in practice macOS and 

80 Windows. macOS exposes no affinity mask at all; Windows does (via 

81 ``GetProcessAffinityMask`` and job-object rate limits) but it is not 

82 implemented here, so a restricted Windows process falls back to the 

83 logical CPU count clamped by ``_MAX_DEFAULT_WORKERS``. 

84 On Linux ``sched_getaffinity`` always succeeds, so an unconstrained box 

85 reports its full logical CPU count rather than ``None``. 

86 

87 Cached for the process lifetime: the CPU allocation is static in practice 

88 (affinity is normally set once at startup, and a live quota change via 

89 ``docker update`` or an in-place pod resize is rare), and this is on the 

90 ``read_csv`` hot path. 

91 """ 

92 limits = [] 

93 if hasattr(os, "sched_getaffinity"): 

94 try: 

95 limits.append(len(os.sched_getaffinity(0))) 

96 except OSError: 

97 pass 

98 quota = _cgroup_cpu_quota() 

99 if quota is not None: 

100 limits.append(max(1, int(quota))) 

101 return min(limits) if limits else None