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

1"""Whether the bundled engine can decode a remote GGUF, answered by gguf-parser.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6import re 

7import subprocess 

8 

9from lilbee.catalog.compat import UnsupportedQuantError 

10from lilbee.providers.base import ProviderError 

11from lilbee.providers.fleet.binary import resolve_gguf_parser 

12from lilbee.providers.fleet.proc import run_bounded 

13 

14log = logging.getLogger(__name__) 

15 

16_FLAG_HF_REPO = "--hf-repo" 

17_FLAG_HF_FILE = "--hf-file" 

18# Named for what it carries: a constant spelled with "token" trips the 

19# hardcoded-credential lint. 

20_FLAG_AUTH = "--token" 

21_FLAG_JSON = "--json" 

22_FLAG_SKIP_ESTIMATE = "--skip-estimate" 

23 

24_PARSER_LABEL = "gguf-parser" 

25 

26# A non-zero exit means the parser could not read the file as a model, which is 

27# the engine's verdict, unless it never got to read it. Listing what a failure to 

28# reach the file looks like is narrower and ages better than listing every way a 

29# file can be rejected: a new rejection reason should count as a verdict, and a 

30# new transport error should not. 

31_ACCESS_FAILURE_MARKERS = ( 

32 "open http file:", 

33 "no such host", 

34 "status code 4", 

35 "status code 5", 

36 "connection refused", 

37 "context deadline exceeded", 

38 "no such file or directory", 

39) 

40_QUANT_TYPE_RE = re.compile(r"GGMLType\(\d+\)") 

41 

42# Header range-reads only, but a large tokenizer pushes the tensor table past 

43# 10 MB, so the ceiling is generous next to the download it guards. 

44_PROBE_TIMEOUT_S = 120 

45_PROBE_KILL_WAIT_S = 5.0 

46 

47 

48def _run_parser(argv: list[str]) -> tuple[str, int] | None: 

49 """(merged output, exit code) from one parser run, or None if it would not run.""" 

50 try: 

51 return run_bounded( 

52 argv, 

53 timeout_s=_PROBE_TIMEOUT_S, 

54 kill_wait_s=_PROBE_KILL_WAIT_S, 

55 merge_stderr=True, 

56 label=_PARSER_LABEL, 

57 ) 

58 except (OSError, subprocess.SubprocessError) as exc: 

59 log.debug("Loadability probe did not run: %s", exc) 

60 return None 

61 

62 

63def _named_quant(detail: str) -> str: 

64 """The quantization gguf-parser named in *detail*, or the reason it gave.""" 

65 match = _QUANT_TYPE_RE.search(detail) 

66 if match: 

67 return match.group(0) 

68 return detail.split(": ", 1)[-1] if ": " in detail else detail 

69 

70 

71def assert_engine_can_load(hf_repo: str, filename: str, token: str | None = None) -> None: 

72 """Raise :class:`UnsupportedQuantError` when the engine cannot decode *filename*. 

73 

74 gguf-parser is the engine's own reader, built from the same pin as 

75 llama-server, so the types it accepts are the binary's rather than a table 

76 lilbee would have to keep in step. It range-reads the header and downloads 

77 none of the weights. 

78 

79 Anything that is not a tensor failure (a 404, DNS, an expired token) is not a 

80 verdict: the probe stays quiet and lets the download report the real problem, 

81 which is how the architecture probe already degrades. An install without the 

82 engine extra has no parser to ask, and must still be able to pull. 

83 """ 

84 try: 

85 parser = resolve_gguf_parser() 

86 except ProviderError as exc: 

87 log.debug("No gguf-parser to check %s/%s against: %s", hf_repo, filename, exc) 

88 return 

89 argv = [ 

90 str(parser), 

91 _FLAG_HF_REPO, 

92 hf_repo, 

93 _FLAG_HF_FILE, 

94 filename, 

95 _FLAG_JSON, 

96 _FLAG_SKIP_ESTIMATE, 

97 ] 

98 if token: 

99 argv += [_FLAG_AUTH, token] 

100 first = _run_parser(argv) 

101 if first is None: 

102 return 

103 output, returncode = first 

104 if returncode == 0: 

105 return 

106 detail = output.strip() 

107 if any(marker in detail for marker in _ACCESS_FAILURE_MARKERS): 

108 log.debug("Loadability probe could not reach %s/%s: %s", hf_repo, filename, detail) 

109 return 

110 # A header runs to megabytes and a read cut short partway through reports as 

111 # a parse failure, indistinguishable by its wording from a real refusal. A 

112 # verdict is a property of the file and repeats; a truncated read does not. 

113 second = _run_parser(argv) 

114 if second is None or second[1] == 0: 

115 log.debug("Loadability probe failed once for %s/%s: %s", hf_repo, filename, detail) 

116 return 

117 raise UnsupportedQuantError(f"{hf_repo}/{filename}", _named_quant(second[0].strip()))