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 

64 def to_state(self) -> dict: 

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

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

67 return { 

68 "role": self.role.value, 

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

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

71 "model": self.model, 

72 "token_cap": self.token_cap, 

73 "weights_bytes": self.weights_bytes, 

74 "slots": self.slots, 

75 "ctx": self.ctx, 

76 "built_ctx_target": self.built_ctx_target, 

77 "built_slots_target": self.built_slots_target, 

78 "replica": self.replica, 

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

80 "est_vram_bytes": self.est_vram_bytes, 

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

82 "est_unreported_bytes": self.est_unreported_bytes, 

83 } 

84 

85 @classmethod 

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

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

88 raw_mode = payload.get("rerank_mode") 

89 return cls( 

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

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

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

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

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

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

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

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

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

99 est_vram_by_device={ 

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

101 }, 

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

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

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

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

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

107 ) 

108 

109 @property 

110 def binary(self) -> str: 

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

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

113 

114 @property 

115 def model_id(self) -> str: 

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

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