Coverage for src/lilbee/providers/fleet/swap_config.py: 100%

1"""Generate one swap group's llama-swap config. 

2 

3Each group runs behind its own llama-swap process with its own config, so a 

4reload of one group never touches another's loaded servers. See 

5docs/architecture.md (llama-swap) for the supervisor/proxy design. 

6""" 

7 

8from __future__ import annotations 

9 

10import json 

11import shlex 

12import subprocess 

13import sys 

14from typing import TYPE_CHECKING 

15 

16from lilbee.data.offload import max_workers 

17from lilbee.providers.fleet.readback import engine_log_env, writes_engine_log 

18from lilbee.providers.roles import WorkerRole 

19 

20if TYPE_CHECKING: 

21 from collections.abc import Mapping 

22 from pathlib import Path 

23 

24 from lilbee.providers.fleet.launch import InstanceLaunch 

25 

26# One group holds this process's members. Swap disabled keeps them co-resident (a 

27# role's replicas); enabled makes llama-swap evict one to load another. 

28_GROUP_NAME = "lilbee" 

29# Cold-load ceiling floor per role; the member's weights scale it up from here. 

30# An embedder's weights stream in seconds, so its floor covers only process start 

31# and backend init: waiting out a chat-sized ten minutes on a crash-looping or 

32# wedged embed server holds every ingest caller for that long. 

33_COLD_LOAD_FLOOR_S: dict[WorkerRole, int] = { 

34 WorkerRole.CHAT: 600, 

35 WorkerRole.VISION: 600, 

36 WorkerRole.RERANK: 600, 

37 WorkerRole.EMBED: 120, 

38} 

39# Conservative cold-load disk rate; a slow network volume streams well under this. 

40_COLD_LOAD_BYTES_PER_S = 150 * 1024 * 1024 

41_LOG_LEVEL = "info" 

42# Matches the --host every server argv binds (adapters); "localhost" would have 

43# llama-swap dial [::1] first, where another process could hold the same port. 

44_PROXY_URL_TEMPLATE = "http://127.0.0.1:{port}" 

45# Shared with the swap manager, whose orphan-server sweep matches this flag's 

46# value in survivor cmdlines. 

47PORT_FLAG = "--port" 

48 

49# llama-swap refuses a request with HTTP 429 once this many are in flight for 

50# one model (its default is 10), and a request held through a cold load counts. 

51# One engine serves every lilbee process the user runs (the TUI, serve with its 

52# HTTP clients, MCP, an ingest run), and llama-server queues requests past its 

53# slots by itself, so the limit only has to stop runaway load. It covers this 

54# many processes, each at its own ceiling: the larger of the slots and the 

55# ingest pool, whose every thread can hold one embed or OCR request. 

56_CONCURRENCY_PROCESSES = 4 

57# Part of the engine pin, so a peer never binds an engine whose config lacks a 

58# key it relies on. 1: every model carries concurrencyLimit. 

59SWAP_CONFIG_REVISION = 1 

60 

61# llama-swap config keys. 

62_KEY_HEALTH_TIMEOUT = "healthCheckTimeout" 

63_KEY_LOG_LEVEL = "logLevel" 

64_KEY_MODELS = "models" 

65_KEY_CMD = "cmd" 

66_KEY_PROXY = "proxy" 

67_KEY_TTL = "ttl" 

68_KEY_CONCURRENCY_LIMIT = "concurrencyLimit" 

69_KEY_ENV = "env" 

70_KEY_GROUPS = "groups" 

71_KEY_SWAP = "swap" 

72_KEY_EXCLUSIVE = "exclusive" 

73_KEY_PERSISTENT = "persistent" 

74_KEY_MEMBERS = "members" 

75 

76 

77def build_swap_config( 

78 launches: list[InstanceLaunch], 

79 member_ports: Mapping[str, int], 

80 *, 

81 swap: bool = False, 

82 ttl_seconds: int = 0, 

83 engine_log_dir: Path | None = None, 

84) -> str: 

85 """Render a llama-swap config (JSON, which is valid YAML) for *launches*. 

86 

87 Each launch becomes a model whose id is its replica model id and whose 

88 command is the llama-server argv plus the explicit port from *member_ports*; 

89 one group holds them behind this group's proxy endpoint. ``swap`` makes the 

90 members evict each other on load, so only one is resident at a time; the 

91 default keeps them co-resident. ``ttl_seconds`` is llama-swap's idle unload 

92 timer per member; 0 keeps weights loaded forever. Ports are allocated fresh per start (never 

93 llama-swap's fixed ``startPort`` range) so a previous instance's lingering 

94 server can't collide with the new fleet's bind. 

95 

96 ``engine_log_dir`` points each engine at its own log there, at the verbosity 

97 that reports what it allocated. llama-swap forwards none of the upstream's 

98 output, so for an engine without ``GET /memory`` that file is the only place 

99 those numbers exist, and reading them back is what tells the planner whether 

100 its estimate held (:mod:`lilbee.providers.fleet.readback`). A launch carrying 

101 ``--memory`` serves the same numbers over HTTP instead and gets no log env: 

102 the trace-level file is exactly what the endpoint retires. 

103 """ 

104 models: dict[str, object] = {} 

105 for launch in launches: 

106 port = member_ports[launch.model_id] 

107 entry: dict[str, object] = { 

108 _KEY_CMD: _command_line(launch.argv, port), 

109 _KEY_PROXY: _PROXY_URL_TEMPLATE.format(port=port), 

110 _KEY_TTL: ttl_seconds, 

111 _KEY_CONCURRENCY_LIMIT: _concurrency_limit(launch), 

112 } 

113 env = dict(launch.env_overrides) 

114 if engine_log_dir is not None and writes_engine_log(launch.argv): 

115 env.update(engine_log_env(engine_log_dir, launch.model_id)) 

116 if env: 

117 entry[_KEY_ENV] = [f"{key}={value}" for key, value in env.items()] 

118 models[launch.model_id] = entry 

119 config: dict[str, object] = { 

120 _KEY_HEALTH_TIMEOUT: _health_check_timeout_s(launches), 

121 _KEY_LOG_LEVEL: _LOG_LEVEL, 

122 _KEY_MODELS: models, 

123 _KEY_GROUPS: { 

124 _GROUP_NAME: { 

125 _KEY_SWAP: swap, 

126 _KEY_EXCLUSIVE: False, 

127 _KEY_PERSISTENT: True, 

128 _KEY_MEMBERS: [launch.model_id for launch in launches], 

129 } 

130 }, 

131 } 

132 return json.dumps(config, indent=2) 

133 

134 

135def _concurrency_limit(launch: InstanceLaunch) -> int: 

136 """Requests llama-swap lets in flight for *launch* before it answers 429.""" 

137 return _CONCURRENCY_PROCESSES * max(launch.slots, max_workers()) 

138 

139 

140def cold_load_timeout_s(weights_bytes: int, role: WorkerRole) -> int: 

141 """Cold-load ceiling for one member's weights at a conservative disk rate, floored per role. 

142 

143 The single source of the scaling formula: llama-swap's health-check timeout 

144 and the provider's per-client request timeout both derive from it, so a model 

145 whose load llama-swap would wait out can never time out the client first. 

146 """ 

147 return max(_COLD_LOAD_FLOOR_S[role], weights_bytes // _COLD_LOAD_BYTES_PER_S) 

148 

149 

150def _health_check_timeout_s(launches: list[InstanceLaunch]) -> int: 

151 """Slowest member's cold-load ceiling; the timeout is proxy-global in 

152 llama-swap, so the slowest possible load sets it.""" 

153 return max( 

154 (cold_load_timeout_s(launch.weights_bytes, launch.role) for launch in launches), 

155 default=max(_COLD_LOAD_FLOOR_S.values()), 

156 ) 

157 

158 

159def _command_line(argv: list[str], port: int) -> str: 

160 """Shell command for a member: the role argv plus its explicit port. 

161 

162 Quoting must match how llama-swap splits the command back into argv: MS 

163 rules on Windows (POSIX single quotes would stay literal in the paths and 

164 the spawn fails with "file does not exist"), POSIX everywhere else. 

165 """ 

166 if sys.platform == "win32": 

167 rendered = subprocess.list2cmdline(argv) 

168 else: 

169 rendered = shlex.join(argv) 

170 return f"{rendered} {PORT_FLAG} {port}"