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 

12 

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

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

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

16 

17 

18@dataclass 

19class InstanceLaunch: 

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

21 

22 role: WorkerRole 

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

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

25 model: str 

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

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

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

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

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

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

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

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

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

35 built_ctx_target: int = 0 

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

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

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

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

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

41 est_vram_bytes: int = 0 

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

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

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

45 # overruns a card. 

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

47 est_unreported_bytes: int = 0 

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

49 

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

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

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

53 """ 

54 

55 def to_state(self) -> dict: 

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

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

58 return { 

59 "role": self.role.value, 

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

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

62 "model": self.model, 

63 "token_cap": self.token_cap, 

64 "weights_bytes": self.weights_bytes, 

65 "slots": self.slots, 

66 "ctx": self.ctx, 

67 "built_ctx_target": self.built_ctx_target, 

68 "replica": self.replica, 

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

70 "est_vram_bytes": self.est_vram_bytes, 

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

72 "est_unreported_bytes": self.est_unreported_bytes, 

73 } 

74 

75 @classmethod 

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

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

78 raw_mode = payload.get("rerank_mode") 

79 return cls( 

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

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

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

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

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

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

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

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

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

89 est_vram_by_device={ 

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

91 }, 

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

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

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

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

96 ) 

97 

98 @property 

99 def model_id(self) -> str: 

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

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