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 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 that file is the only place those numbers exist, and reading them 

76 back is what tells the planner whether its estimate held 

77 (:mod:`lilbee.providers.fleet.readback`). 

78 """ 

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

80 for launch in launches: 

81 port = member_ports[launch.model_id] 

82 entry: dict[str, object] = { 

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

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

85 _KEY_TTL: ttl_seconds, 

86 } 

87 env = dict(launch.env_overrides) 

88 if engine_log_dir is not None: 

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

90 if env: 

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

92 models[launch.model_id] = entry 

93 config: dict[str, object] = { 

94 _KEY_HEALTH_TIMEOUT: _health_check_timeout_s(launches), 

95 _KEY_LOG_LEVEL: _LOG_LEVEL, 

96 _KEY_MODELS: models, 

97 _KEY_GROUPS: { 

98 _GROUP_NAME: { 

99 _KEY_SWAP: swap, 

100 _KEY_EXCLUSIVE: False, 

101 _KEY_PERSISTENT: True, 

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

103 } 

104 }, 

105 } 

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

107 

108 

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

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

111 

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

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

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

115 """ 

116 return max(_HEALTH_CHECK_TIMEOUT_FLOOR_S, weights_bytes // _COLD_LOAD_BYTES_PER_S) 

117 

118 

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

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

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

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

123 return cold_load_timeout_s(heaviest) 

124 

125 

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

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

128 

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

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

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

132 """ 

133 if sys.platform == "win32": 

134 rendered = subprocess.list2cmdline(argv) 

135 else: 

136 rendered = shlex.join(argv) 

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