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 

25from lilbee.runtime.shard_process import shard_process 

26 

27log = logging.getLogger(__name__) 

28 

29_ENV_VAR = "LILBEE_CPU_QUOTA" 

30 

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

32 

33 

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

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

36 

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

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

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

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

41 per caller, so it is left to them. 

42 """ 

43 try: 

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

45 except OSError: 

46 v2 = [] 

47 if v2: 

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

49 return None 

50 try: 

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

52 except (ValueError, IndexError): 

53 return None 

54 return quota / period if period > 0 else None 

55 try: 

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

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

58 except (OSError, ValueError): 

59 return None 

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

61 return None 

62 return quota / period 

63 

64 

65def _cpu_budget() -> float: 

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

67 

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

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

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

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

72 

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

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

75 """ 

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

77 if hasattr(os, "sched_getaffinity"): 

78 with contextlib.suppress(OSError): 

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

80 quota = _cgroup_cpu_quota() 

81 if quota is not None: 

82 limits.append(quota) 

83 return min(limits) 

84 

85 

86def available_cpu_count() -> int: 

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

88 

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

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

91 """ 

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

93 

94 

95def cpu_quota() -> int: 

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

97 

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

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

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

101 """ 

102 if shard_process.cpu_quota is not None: 

103 return shard_process.cpu_quota 

104 override = os.environ.get(_ENV_VAR) 

105 if override is not None: 

106 try: 

107 value = int(override) 

108 if value > 0: 

109 return value 

110 except ValueError: 

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

112 log.warning( 

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

114 _ENV_VAR, 

115 override, 

116 ) 

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

118 

119 

120def engine_thread_count() -> int | None: 

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

122 

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

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

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

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

127 count is the better informed one and stays. 

128 

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

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

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

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

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

134 

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

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

137 """ 

138 budget = _cpu_budget() 

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

140 return None 

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