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.providers.fleet.readback import MEMORY_FLAG, engine_log_env 

17 

18if TYPE_CHECKING: 

19 from collections.abc import Mapping 

20 from pathlib import Path 

21 

22 from lilbee.providers.fleet.launch import InstanceLaunch 

23 

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

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

26_GROUP_NAME = "lilbee" 

27# Cold-load ceiling floor; the heaviest member's weights scale it up from here. 

28_HEALTH_CHECK_TIMEOUT_FLOOR_S = 600 

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

30_COLD_LOAD_BYTES_PER_S = 150 * 1024 * 1024 

31_LOG_LEVEL = "info" 

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

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

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

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

36# value in survivor cmdlines. 

37PORT_FLAG = "--port" 

38 

39# llama-swap config keys. 

40_KEY_HEALTH_TIMEOUT = "healthCheckTimeout" 

41_KEY_LOG_LEVEL = "logLevel" 

42_KEY_MODELS = "models" 

43_KEY_CMD = "cmd" 

44_KEY_PROXY = "proxy" 

45_KEY_TTL = "ttl" 

46_KEY_ENV = "env" 

47_KEY_GROUPS = "groups" 

48_KEY_SWAP = "swap" 

49_KEY_EXCLUSIVE = "exclusive" 

50_KEY_PERSISTENT = "persistent" 

51_KEY_MEMBERS = "members" 

52 

53 

54def build_swap_config( 

55 launches: list[InstanceLaunch], 

56 member_ports: Mapping[str, int], 

57 *, 

58 swap: bool = False, 

59 ttl_seconds: int = 0, 

60 engine_log_dir: Path | None = None, 

61) -> str: 

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

63 

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

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

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

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

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

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

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

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

72 

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

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

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

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

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

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

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

80 """ 

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

82 for launch in launches: 

83 port = member_ports[launch.model_id] 

84 entry: dict[str, object] = { 

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

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

87 _KEY_TTL: ttl_seconds, 

88 } 

89 env = dict(launch.env_overrides) 

90 if engine_log_dir is not None and MEMORY_FLAG not in launch.argv: 

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

92 if env: 

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

94 models[launch.model_id] = entry 

95 config: dict[str, object] = { 

96 _KEY_HEALTH_TIMEOUT: _health_check_timeout_s(launches), 

97 _KEY_LOG_LEVEL: _LOG_LEVEL, 

98 _KEY_MODELS: models, 

99 _KEY_GROUPS: { 

100 _GROUP_NAME: { 

101 _KEY_SWAP: swap, 

102 _KEY_EXCLUSIVE: False, 

103 _KEY_PERSISTENT: True, 

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

105 } 

106 }, 

107 } 

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

109 

110 

111def cold_load_timeout_s(weights_bytes: int) -> int: 

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

113 

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

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

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

117 """ 

118 return max(_HEALTH_CHECK_TIMEOUT_FLOOR_S, weights_bytes // _COLD_LOAD_BYTES_PER_S) 

119 

120 

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

122 """Cold-load ceiling of the heaviest member; the timeout is proxy-global in 

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

124 heaviest = max((launch.weights_bytes for launch in launches), default=0) 

125 return cold_load_timeout_s(heaviest) 

126 

127 

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

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

130 

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

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

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

134 """ 

135 if sys.platform == "win32": 

136 rendered = subprocess.list2cmdline(argv) 

137 else: 

138 rendered = shlex.join(argv) 

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