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

1"""The launch spec for one llama-server instance, shared by planning and llama-swap.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass, field 

6 

7from lilbee.providers.roles import RerankMode, WorkerRole 

8 

9# Separator between a role and its replica index in a llama-swap model id. 

10_REPLICA_SEP = "-" 

11# Stand-in for the binary of a record whose argv arrived empty (a foreign state file). 

12_UNKNOWN_BINARY = "an unrecorded binary" 

13 

14 

15def role_model_prefix(role: WorkerRole) -> str: 

16 """Prefix shared by every replica model id of *role* (``<role>-``).""" 

17 return f"{role.value}{_REPLICA_SEP}" 

18 

19 

20@dataclass 

21class InstanceLaunch: 

22 """Everything needed to run one llama-server, minus the port (claimed at spawn).""" 

23 

24 role: WorkerRole 

25 argv: list[str] # llama-server command WITHOUT --port; the runner appends it 

26 env_overrides: dict[str, str] # backend-specific device-pinning env 

27 model: str 

28 token_cap: int | None = None # per-slot ctx for embed/rerank input truncation 

29 weights_bytes: int = 0 # model file size on disk; scales the cold-load timeout 

30 slots: int = 1 # --parallel continuous-batching slots; chat concurrency capacity 

31 ctx: int = 0 # per-slot context the server runs with; what a client should fit to 

32 # The chat ctx target the builder planned against (num_ctx pin else 

33 # chat_n_ctx_target); 0 for non-chat roles and pre-field records. A 

34 # co-tenant whose target this covers adopts the engine even when the 

35 # served window is smaller: the same planner aiming at least as high 

36 # already achieved this window, so a rebuild cannot beat it. 

37 built_ctx_target: int = 0 

38 # The vision slot ceiling the builder fitted against (vision_ocr_concurrency); 

39 # 0 for other roles and pre-field records. 

40 built_slots_target: int = 0 

41 replica: int = 0 # index within the role's data-parallel pool (0 = single server) 

42 rerank_mode: RerankMode | None = None # set only for RERANK; picks the client scoring path 

43 # GPU bytes placement charged this instance. Carried so the engine's own 

44 # startup report can be checked against it once the server is up; 0 for a 

45 # model the estimator could not size, where there is nothing to compare. 

46 est_vram_bytes: int = 0 

47 # What placement charged each card this instance runs on, keyed by the name 

48 # the engine prints for it. The scalar above cannot distinguish a split that 

49 # landed 50/50 from one that landed 80/20, and the second is the one that 

50 # overruns a card. 

51 est_vram_by_device: dict[str, int] = field(default_factory=dict) 

52 est_unreported_bytes: int = 0 

53 """Estimated bytes the engine allocates but never reports in its buffer lines. 

54 

55 A vision projector's weights are the case: llama.cpp allocates them without 

56 emitting a "buffer size" line, so the readback total is short by exactly this 

57 much and the self-check would warn on a load that was sized correctly. 

58 

59 Log-mode readback only. An engine serving GET /memory reports the projector 

60 per device in its mmproj field, so that path compares the full estimate and 

61 ignores this. 

62 """ 

63 settings: str = "" 

64 est_host_bytes: int = 0 

65 """Host bytes the estimate expects chat to hold when the engine fits its layers. 

66 

67 The readback says the engine left chat weights in system memory when the 

68 host holds more than this. 0 where the engine's fit does not run. 

69 """ 

70 

71 def to_state(self) -> dict: 

72 """JSON-safe form written into every engine state file so a guest lilbee 

73 can read the serving contract and bind.""" 

74 return { 

75 "role": self.role.value, 

76 "argv": list(self.argv), 

77 "env_overrides": dict(self.env_overrides), 

78 "model": self.model, 

79 "token_cap": self.token_cap, 

80 "weights_bytes": self.weights_bytes, 

81 "slots": self.slots, 

82 "ctx": self.ctx, 

83 "built_ctx_target": self.built_ctx_target, 

84 "built_slots_target": self.built_slots_target, 

85 "replica": self.replica, 

86 "rerank_mode": self.rerank_mode.value if self.rerank_mode else None, 

87 "est_vram_bytes": self.est_vram_bytes, 

88 "est_vram_by_device": dict(self.est_vram_by_device), 

89 "est_unreported_bytes": self.est_unreported_bytes, 

90 "est_host_bytes": self.est_host_bytes, 

91 "settings": self.settings, 

92 } 

93 

94 @classmethod 

95 def from_state(cls, payload: dict) -> InstanceLaunch: 

96 """Rebuild a launch from :meth:`to_state` output; raises on a foreign shape.""" 

97 raw_mode = payload.get("rerank_mode") 

98 return cls( 

99 role=WorkerRole(payload["role"]), 

100 argv=list(payload["argv"]), 

101 env_overrides=dict(payload.get("env_overrides") or {}), 

102 model=str(payload["model"]), 

103 token_cap=payload.get("token_cap"), 

104 weights_bytes=int(payload.get("weights_bytes") or 0), 

105 slots=int(payload.get("slots") or 1), 

106 est_vram_bytes=int(payload.get("est_vram_bytes") or 0), 

107 est_unreported_bytes=int(payload.get("est_unreported_bytes") or 0), 

108 est_host_bytes=int(payload.get("est_host_bytes") or 0), 

109 est_vram_by_device={ 

110 str(k): int(v) for k, v in (payload.get("est_vram_by_device") or {}).items() 

111 }, 

112 ctx=int(payload.get("ctx") or 0), 

113 built_ctx_target=int(payload.get("built_ctx_target") or 0), 

114 built_slots_target=int(payload.get("built_slots_target") or 0), 

115 replica=int(payload.get("replica") or 0), 

116 rerank_mode=RerankMode(raw_mode) if raw_mode else None, 

117 settings=str(payload.get("settings") or ""), 

118 ) 

119 

120 @property 

121 def binary(self) -> str: 

122 """The llama-server this instance runs: the argv's first word.""" 

123 return next(iter(self.argv), _UNKNOWN_BINARY) 

124 

125 @property 

126 def model_id(self) -> str: 

127 """The llama-swap model id for this instance: ``<role>-<replica>``.""" 

128 return f"{role_model_prefix(self.role)}{self.replica}"