Coverage for src/lilbee/catalog/header_probe.py: 100%

1"""Range-GET a GGUF blob's header and read the fields that identify the file. 

2 

3The fields are read by walking the metadata KV table directly (see 

4``_parse_header``) rather than via gguf-py's ``GGUFReader``, which needs the whole 

5KV table -- including a model's multi-megabyte tokenizer arrays -- present, and 

6so chokes on a Range-GET-truncated header. 

7 

8The bundled gguf-parser reports the same two fields and is the authority on 

9whether the engine can load a file, but it always reads through to the tensor 

10table, which sits past those tokenizer arrays. Measured on the same files it 

11costs 6.7s against this probe's 0.7s for a small model. Selection asks about 

12every candidate and needs only ``general.type``, the second key in the table, so 

13it reads a 64 KB window here; the load verdict is left to gguf-parser, once, in 

14``providers.fleet.loadability``. 

15""" 

16 

17from __future__ import annotations 

18 

19import logging 

20import struct 

21from dataclasses import dataclass 

22from http import HTTPStatus 

23 

24import httpx 

25 

26log = logging.getLogger(__name__) 

27 

28GGUF_HEADER_PROBE_BYTES = 65536 

29GGUF_MAGIC = b"GGUF" 

30GGUF_ARCH_KEY = "general.architecture" 

31GGUF_TYPE_KEY = "general.type" 

32 

33# ``general.type`` values that name something other than loadable model weights. 

34# A file that omits the key predates it and counts as a model, so the check only 

35# rejects a file whose header says outright that it is not one. 

36NON_MODEL_GGUF_TYPES = frozenset({"mmproj", "adapter"}) 

37 

38_PROBE_TIMEOUT_S = 10.0 

39 

40 

41@dataclass(frozen=True) 

42class GgufHeader: 

43 """What a GGUF file says it is. Empty fields mean the header was unreadable.""" 

44 

45 architecture: str = "" 

46 file_type: str = "" 

47 

48 @property 

49 def is_model(self) -> bool: 

50 """True unless the header identifies the file as a projector or adapter.""" 

51 return self.file_type not in NON_MODEL_GGUF_TYPES 

52 

53 

54def probe_header(blob_url: str) -> GgufHeader: 

55 """Read a GGUF blob's identifying header fields, empty on any failure.""" 

56 try: 

57 headers = {"Range": f"bytes=0-{GGUF_HEADER_PROBE_BYTES - 1}"} 

58 # follow_redirects: a HuggingFace ``/resolve/`` URL for an LFS-backed GGUF 

59 # 302-redirects to the CDN; without following it the probe reads the tiny 

60 # redirect body, the GGUF magic check fails, and every arch verdict is 

61 # silently UNKNOWN (so the unsupported-arch guard never fires). 

62 resp = httpx.get(blob_url, headers=headers, timeout=_PROBE_TIMEOUT_S, follow_redirects=True) 

63 if resp.status_code >= HTTPStatus.BAD_REQUEST: 

64 return GgufHeader() 

65 return _parse_header(resp.content) 

66 except httpx.HTTPError as exc: 

67 log.debug("GGUF header probe failed for %s: %s", blob_url, exc) 

68 return GgufHeader() 

69 

70 

71# GGUF metadata value-type tags (gguf spec) and the fixed byte sizes of the 

72# scalar ones. ARRAY/STRING carry their own length prefixes. 

73_GGUF_TYPE_STRING = 8 

74_GGUF_TYPE_ARRAY = 9 

75_GGUF_SCALAR_SIZES = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8} 

76# magic(4) + version(4) + tensor_count(8) + kv_count(8) 

77_GGUF_HEADER_FIXED = 24 

78 

79 

80class _TruncatedHeaderError(Exception): 

81 """The probe window ended before the bytes a parse step needed.""" 

82 

83 

84class _HeaderCursor: 

85 """Little-endian reader over the probed GGUF header bytes.""" 

86 

87 def __init__(self, blob: bytes) -> None: 

88 self._blob = blob 

89 self._pos = 0 

90 

91 def take(self, n: int) -> bytes: 

92 end = self._pos + n 

93 if n < 0 or end > len(self._blob): 

94 raise _TruncatedHeaderError 

95 chunk = self._blob[self._pos : end] 

96 self._pos = end 

97 return chunk 

98 

99 def u32(self) -> int: 

100 return int(struct.unpack_from("<I", self.take(4))[0]) 

101 

102 def u64(self) -> int: 

103 return int(struct.unpack_from("<Q", self.take(8))[0]) 

104 

105 def gguf_string(self) -> bytes: 

106 return self.take(self.u64()) 

107 

108 

109def _skip_value(cur: _HeaderCursor, value_type: int) -> None: 

110 """Advance *cur* past one metadata value of *value_type*.""" 

111 size = _GGUF_SCALAR_SIZES.get(value_type) 

112 if size is not None: 

113 cur.take(size) 

114 return 

115 if value_type == _GGUF_TYPE_STRING: 

116 cur.gguf_string() 

117 return 

118 if value_type == _GGUF_TYPE_ARRAY: 

119 elem_type = cur.u32() 

120 count = cur.u64() 

121 elem_size = _GGUF_SCALAR_SIZES.get(elem_type) 

122 if elem_size is not None: 

123 cur.take(elem_size * count) 

124 elif elem_type == _GGUF_TYPE_STRING: 

125 for _ in range(count): 

126 cur.gguf_string() 

127 else: 

128 raise _TruncatedHeaderError # nested/unknown array element: stop parsing 

129 return 

130 raise _TruncatedHeaderError # unknown value type: stop parsing 

131 

132 

133_WANTED_STRING_KEYS: dict[bytes, str] = { 

134 GGUF_ARCH_KEY.encode("utf-8"): "architecture", 

135 GGUF_TYPE_KEY.encode("utf-8"): "file_type", 

136} 

137_GGUF_MAGIC_LEN = 4 

138 

139 

140def _collect_string_fields(cur: _HeaderCursor, kv_count: int) -> dict[str, str]: 

141 """Walk the KV table and return the ``_WANTED_STRING_KEYS`` values it carries. 

142 

143 Keeps what it read when the probe window ends mid-table, so a header whose 

144 wanted keys precede a huge tokenizer array still resolves. 

145 """ 

146 found: dict[str, str] = {} 

147 try: 

148 for _ in range(kv_count): 

149 key = cur.gguf_string() 

150 value_type = cur.u32() 

151 field = _WANTED_STRING_KEYS.get(key) 

152 if field is not None and value_type == _GGUF_TYPE_STRING: 

153 found[field] = cur.gguf_string().decode("utf-8", errors="replace") 

154 if len(found) == len(_WANTED_STRING_KEYS): 

155 break 

156 else: 

157 _skip_value(cur, value_type) 

158 except _TruncatedHeaderError: 

159 pass 

160 return found 

161 

162 

163def _parse_header(blob: bytes) -> GgufHeader: 

164 """Extract the identifying fields from a (possibly truncated) GGUF header. 

165 

166 Walks the metadata KV table directly and stops once every wanted key is read; 

167 GGUF writers emit ``general.architecture`` and ``general.type`` among the 

168 first entries. This deliberately avoids gguf-py's ``GGUFReader``, which parses 

169 the entire KV table up front: a real model's multi-megabyte tokenizer arrays 

170 run past the Range-GET probe window, so ``GGUFReader`` raises on the 

171 truncation before any field can be read. A malformed or too-short header 

172 yields whatever was read before the truncation. 

173 """ 

174 if len(blob) < _GGUF_HEADER_FIXED or blob[:_GGUF_MAGIC_LEN] != GGUF_MAGIC: 

175 return GgufHeader() 

176 # The length guard above covers exactly these reads, so none of them can 

177 # run off the end; only the KV walk that follows can truncate. 

178 cur = _HeaderCursor(blob) 

179 cur.take(_GGUF_MAGIC_LEN) 

180 cur.u32() # version 

181 cur.u64() # tensor_count (the tensor-info table is never read) 

182 return GgufHeader(**_collect_string_fields(cur, cur.u64()))