Coverage for src/lilbee/runtime/cpu.py: 100%

1"""CPU concurrency policy for compute-bound parallelism. 

2 

3Use ``cpu_quota()`` to bound thread pools and asyncio semaphores that 

4schedule CPU-heavy work (chat inference, PDF rasterization, embedding, 

5tokenization). Capping at half of cpu_count keeps the asyncio main 

6thread scheduler share so the TUI stays responsive while a worker 

7storm is in flight. 

8 

9Do NOT use this for HTTP request concurrency: that lives under 

10``cfg.crawl_max_concurrent`` and is governed by remote-side rate 

11limits, not local CPU. 

12 

13``engine_thread_count()`` is the out-of-process counterpart, sizing a 

14spawned llama-server's thread count. 

15""" 

16 

17from __future__ import annotations 

18 

19import contextlib 

20import logging 

21import math 

22import os 

23from pathlib import Path 

24 

25log = logging.getLogger(__name__) 

26 

27_ENV_VAR = "LILBEE_CPU_QUOTA" 

28 

29_CGROUP_ROOT = Path("/sys/fs/cgroup") 

30 

31 

32def _cgroup_cpu_quota(root: Path = _CGROUP_ROOT) -> float | None: 

33 """Cores the CFS quota allows, unrounded, or None when unlimited or unreadable. 

34 

35 cgroup v2 keeps ``<quota> <period>`` (or ``max`` for unlimited) in 

36 ``cpu.max``; v1 splits it across ``cpu.cfs_quota_us`` (-1 = unlimited) and 

37 ``cpu.cfs_period_us``. Quotas are routinely fractional (a rented pod gets 

38 765ms per 100ms period, i.e. 7.65 cores), and which way that rounds differs 

39 per caller, so it is left to them. 

40 """ 

41 try: 

42 v2 = (root / "cpu.max").read_text(encoding="utf-8").split() 

43 except OSError: 

44 v2 = [] 

45 if v2: 

46 if v2[0] == "max": 

47 return None 

48 try: 

49 quota, period = int(v2[0]), int(v2[1]) 

50 except (ValueError, IndexError): 

51 return None 

52 return quota / period if period > 0 else None 

53 try: 

54 quota = int((root / "cpu" / "cpu.cfs_quota_us").read_text(encoding="utf-8")) 

55 period = int((root / "cpu" / "cpu.cfs_period_us").read_text(encoding="utf-8")) 

56 except (OSError, ValueError): 

57 return None 

58 if quota <= 0 or period <= 0: 

59 return None 

60 return quota / period 

61 

62 

63def _cpu_budget() -> float: 

64 """Cores this process may use, unrounded: the tightest of the three signals. 

65 

66 ``os.cpu_count()`` reports the host's cores, which over-reports inside a 

67 CPU-limited container (a rented multi-vCPU box hands a pod a fraction of the 

68 machine). Fold in the process's scheduling affinity and the cgroup CFS quota 

69 so a container-bound run sizes to its real budget rather than the host's. 

70 

71 ``os.process_cpu_count()`` folds these in for us, but it landed in 3.13 and 

72 the project floor is 3.11, so the cgroup read is done by hand here. 

73 """ 

74 limits = [float(os.cpu_count() or 1)] 

75 if hasattr(os, "sched_getaffinity"): 

76 with contextlib.suppress(OSError): 

77 limits.append(float(len(os.sched_getaffinity(0)))) 

78 quota = _cgroup_cpu_quota() 

79 if quota is not None: 

80 limits.append(quota) 

81 return min(limits) 

82 

83 

84def available_cpu_count() -> int: 

85 """Usable CPUs for this process, honoring cgroup quota and CPU affinity. 

86 

87 A fractional quota rounds up: a worker pool of that size is right for work 

88 that blocks as well as computes. Always at least 1. 

89 """ 

90 return max(1, math.ceil(_cpu_budget())) 

91 

92 

93def cpu_quota() -> int: 

94 """Return the CPU concurrency cap; honors ``LILBEE_CPU_QUOTA`` override. 

95 

96 Default is ``max(1, available_cpu_count() // 2)``. The override accepts a 

97 positive integer; non-positive or unparseable values fall back to 

98 the default and a warning is logged once per call. 

99 """ 

100 override = os.environ.get(_ENV_VAR) 

101 if override is not None: 

102 try: 

103 value = int(override) 

104 if value > 0: 

105 return value 

106 except ValueError: 

107 pass # bad override falls through to the warning + default below 

108 log.warning( 

109 "Ignoring %s=%r: must be a positive integer; using default.", 

110 _ENV_VAR, 

111 override, 

112 ) 

113 return max(1, available_cpu_count() // 2) 

114 

115 

116def engine_thread_count() -> int | None: 

117 """Threads for a spawned engine, or None to keep the engine's own default. 

118 

119 llama.cpp sizes its default from host topology (physical cores, minus 

120 efficiency cores) and cannot see a CFS quota or an affinity mask, so inside a 

121 CPU-limited container it starts several times more threads than the quota 

122 admits and generation collapses. Where nothing caps this process, upstream's 

123 count is the better informed one and stays. 

124 

125 The budget rounds DOWN here, unlike :func:`available_cpu_count`. An engine's 

126 threads compute in lockstep across a barrier, so one thread over the quota 

127 gets the whole group throttled and every other thread waits for it. Measured 

128 on a 96-core pod with a 7.65-core quota, one server generating: 11.4 tok/s on 

129 the engine's own count, 60.3 at 8 threads, 63.5 at 7. 

130 

131 Not ``cpu_quota()``: that halves the budget to leave the asyncio scheduler a 

132 share, which an out-of-process engine does not compete with. 

133 """ 

134 budget = _cpu_budget() 

135 if budget >= (os.cpu_count() or 1): 

136 return None 

137 return max(1, math.floor(budget))