diff --git a/application/parser/remote/crawler_markdown.py b/application/parser/remote/crawler_markdown.py index dec25e06..c78da24e 100644 --- a/application/parser/remote/crawler_markdown.py +++ b/application/parser/remote/crawler_markdown.py @@ -9,6 +9,11 @@ from application.parser.schema.base import Document import tldextract import os +# The bundled public-suffix snapshot is enough for domain matching; the +# default extractor would fetch the live list on first use and cache it on +# disk, which is a network round trip the ingest worker should not depend on. +_extract = tldextract.TLDExtract(suffix_list_urls=(), cache_dir=None) + class CrawlerLoader(BaseRemote): def __init__(self, limit=10, allow_subdomains=False): """ @@ -124,7 +129,7 @@ class CrawlerLoader(BaseRemote): return links def _get_base_domain(self, url): - extracted = tldextract.extract(url) + extracted = _extract(url) # Reconstruct the domain as domain.suffix base_domain = f"{extracted.domain}.{extracted.suffix}" return base_domain @@ -141,7 +146,7 @@ class CrawlerLoader(BaseRemote): if not parsed_link.netloc: continue - extracted = tldextract.extract(parsed_link.netloc) + extracted = _extract(parsed_link.netloc) link_base = f"{extracted.domain}.{extracted.suffix}" if self.allow_subdomains: diff --git a/application/parser/tokenization.py b/application/parser/tokenization.py index 88059644..c449a7dc 100644 --- a/application/parser/tokenization.py +++ b/application/parser/tokenization.py @@ -216,12 +216,28 @@ class HuggingFaceCounter(TokenCounter): return _cap_piece_chars(pieces, first, rest) +def _tokenizer_file(repo: str) -> str: + """Path to ``repo``'s ``tokenizer.json``, from the hub cache when present. + + A warmed cache (the Docker image bakes the default models) answers without + touching the network. ``hf_hub_download`` would otherwise revalidate the + revision with a HEAD request on every process start, and stall for the + etag timeout on a host that cannot reach huggingface.co. + """ + from huggingface_hub import hf_hub_download + + try: + return hf_hub_download(repo, "tokenizer.json", local_files_only=True) + except Exception: # noqa: BLE001 -- not cached: fetch it + return hf_hub_download(repo, "tokenizer.json") + + def _load_hf_counter(repo: str) -> Optional[HuggingFaceCounter]: """Load ``repo``'s tokenizer, or ``None`` if it is not reachable.""" try: from tokenizers import Tokenizer - tokenizer = Tokenizer.from_pretrained(repo) + tokenizer = Tokenizer.from_file(_tokenizer_file(repo)) # Repos ship padding and truncation defaults meant for inference # batches. Left on, every count returns the padded width (128 for # mpnet) and the offsets carry (0, 0) entries for the padding, so both diff --git a/application/scripts/prefetch_models.py b/application/scripts/prefetch_models.py index 64c6db7f..029f1e3a 100644 --- a/application/scripts/prefetch_models.py +++ b/application/scripts/prefetch_models.py @@ -1,9 +1,14 @@ -"""Download embedding model artifacts into FastEmbed's cache. +"""Download the model artifacts a fresh container would otherwise fetch. -Run at image build time so a fresh container does not download a model on its -first ingest, and an air-gapped install works at all. Both the legacy and the -current default are baked: an upgraded deployment keeps using mpnet until it -runs ``reembed``, while a new one starts on granite. +Run at image build time so a fresh container does not download on its first +request, and an air-gapped install works at all. Two things are warmed: + +* Embedding models, into FastEmbed's cache. Both the legacy and the current + default are baked: an upgraded deployment keeps using mpnet until it runs + ``reembed``, while a new one starts on granite. +* tiktoken's ``cl100k_base`` encoding, which token accounting uses on every + chat. tiktoken caches it under ``TIKTOKEN_CACHE_DIR`` (a temp dir when + unset), so the image sets that variable and this warms it. Usage:: @@ -29,6 +34,23 @@ logger = logging.getLogger("prefetch_models") #: Fetched when no names are given. DEFAULT_MODELS = (DEFAULT_LEGACY, DEFAULT_NEW_INSTALL) +#: tiktoken encodings the application loads (``application.utils.get_encoding``). +TIKTOKEN_ENCODINGS = ("cl100k_base",) + + +def prefetch_tiktoken(names: Sequence[str] = TIKTOKEN_ENCODINGS) -> List[str]: + """Warm tiktoken's cache for each encoding in ``names``. + + Returns: + The encodings fetched. + """ + import tiktoken + + for name in names: + logger.info("Fetching tiktoken encoding %s", name) + tiktoken.get_encoding(name) + return list(names) + def prefetch(names: Sequence[str], cache_dir: Optional[str] = None) -> List[str]: """Fetch each named model's artifacts. @@ -82,6 +104,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int: names = list(argv) if argv else list(DEFAULT_MODELS) fetched = prefetch(names, os.environ.get("EMBEDDINGS_CACHE_DIR")) logger.info("Cached %d model(s): %s", len(fetched), ", ".join(fetched)) + encodings = prefetch_tiktoken() + logger.info("Cached tiktoken encoding(s): %s", ", ".join(encodings)) return 0 diff --git a/application/scripts/verify_offline.py b/application/scripts/verify_offline.py new file mode 100644 index 00000000..a1e560ce --- /dev/null +++ b/application/scripts/verify_offline.py @@ -0,0 +1,153 @@ +"""Check that an image can serve its defaults without any network access. + +Exercises the code paths a fresh container hits first, the way the +application does: tiktoken token counting, the chunker's tokenizer for each +baked embedding model, and a FastEmbed embed with each. Run it inside the +image with networking disabled; every check must pass with zero requests:: + + docker run --rm --network none arc53/docsgpt:latest \\ + python -m application.scripts.verify_offline + +Exit status is non-zero on the first failure. Models to check default to +the prefetch defaults; pass registry names to check a different set. +""" + +from __future__ import annotations + +import logging +import socket +import sys +import time +from typing import Callable, List, Optional, Sequence + +from application.core.optional_deps import is_available +from application.scripts.prefetch_models import DEFAULT_MODELS, TIKTOKEN_ENCODINGS +from application.vectorstore.model_registry import resolve + +logger = logging.getLogger("verify_offline") + + +def _network_reachable(host: str = "huggingface.co", port: int = 443) -> bool: + try: + socket.create_connection((host, port), timeout=2).close() + return True + except OSError: + return False + + +def _check(name: str, fn: Callable[[], object]) -> bool: + started = time.time() + try: + detail = fn() + except Exception as exc: # noqa: BLE001 -- report every failure the same way + print(f"FAIL {name}: {type(exc).__name__}: {exc}") + return False + print(f"ok {name}: {detail} ({time.time() - started:.2f}s)") + return True + + +def verify(models: Sequence[str]) -> bool: + """Run every check; return whether all passed.""" + ok = True + + def tiktoken_check(encoding: str) -> Callable[[], object]: + def run() -> object: + import tiktoken + + return f"{len(tiktoken.get_encoding(encoding).encode('hello world'))} tokens" + + return run + + for encoding in TIKTOKEN_ENCODINGS: + ok &= _check(f"tiktoken {encoding}", tiktoken_check(encoding)) + + for name in models: + spec = resolve(name) + if spec is None or spec.provider != "fastembed": + print(f"skip {name}: not a local model") + continue + + def tokenizer_check(model_name: str = name) -> object: + from application.parser.tokenization import get_token_counter + + counter = get_token_counter(model_name) + if counter.name == "cl100k_base": + raise RuntimeError("tokenizer missing from the cache; chunking fell back to cl100k") + return f"{counter.name}, {counter.count('The quick brown fox')} tokens" + + def embed_check(model_name: str = name) -> object: + from application.vectorstore.embeddings_local import EmbeddingsWrapper + + vector = EmbeddingsWrapper(model_name).embed_query("hello") + return f"dimension {len(vector)}" + + ok &= _check(f"tokenizer {name}", tokenizer_check) + ok &= _check(f"embeddings {name}", embed_check) + + if is_available("docling"): + ok &= _check("docling PDF conversion", _docling_check) + else: + print("skip docling: not installed (slim image)") + + return ok + + +# A one-page PDF with a single text run; enough for the layout model to have +# something to look at. +_TINY_PDF = ( + b"%PDF-1.4\n" + b"1 0 obj<>endobj\n" + b"2 0 obj<>endobj\n" + b"3 0 obj<>>>>>endobj\n" + b"4 0 obj<>stream\n" + b"BT /F1 18 Tf 20 100 Td (Offline verification page) Tj ET\n" + b"endstream\nendobj\n" + b"5 0 obj<>endobj\n" + b"trailer<>\n%%EOF\n" +) + + +def _docling_check() -> object: + """Convert a tiny PDF through docling; its models must come from DOCLING_ARTIFACTS_PATH.""" + import os + import tempfile + + from docling.datamodel.base_models import InputFormat + from docling.datamodel.pipeline_options import PdfPipelineOptions + from docling.document_converter import DocumentConverter, PdfFormatOption + + from application.parser.file.docling_parser import _apply_inference_settings + + # Same global docling settings the parser applies: torch.compile stays off + # unless DOCLING_COMPILE_TORCH_MODELS asks for it (it needs a C++ toolchain). + _apply_inference_settings() + artifacts = os.environ.get("DOCLING_ARTIFACTS_PATH") + options = PdfPipelineOptions(artifacts_path=artifacts, do_ocr=False, do_table_structure=True) + converter = DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=options)}) + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as handle: + handle.write(_TINY_PDF) + path = handle.name + try: + text = converter.convert(path).document.export_to_markdown() + finally: + os.unlink(path) + if "Offline verification" not in text: + raise RuntimeError(f"unexpected conversion output: {text[:80]!r}") + return f"models from {artifacts or 'default cache'}, {len(text)} chars" + + +def main(argv: Optional[Sequence[str]] = None) -> int: + logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(message)s") + models: List[str] = list(argv) if argv else list(DEFAULT_MODELS) + if _network_reachable(): + print("note network is reachable; run with --network none to prove the offline path") + else: + print("note network unreachable, as intended") + passed = verify(models) + print("VERIFY OFFLINE: " + ("PASS" if passed else "FAIL")) + return 0 if passed else 1 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/application/vectorstore/embeddings_local.py b/application/vectorstore/embeddings_local.py index f35854c5..060bc58b 100644 --- a/application/vectorstore/embeddings_local.py +++ b/application/vectorstore/embeddings_local.py @@ -104,7 +104,11 @@ def _read_repo_json(repo: str, filename: str) -> Optional[dict]: try: from huggingface_hub import hf_hub_download - with open(hf_hub_download(repo_id=repo, filename=filename), encoding="utf-8") as handle: + try: + path = hf_hub_download(repo_id=repo, filename=filename, local_files_only=True) + except Exception: # noqa: BLE001 -- not cached: fetch it + path = hf_hub_download(repo_id=repo, filename=filename) + with open(path, encoding="utf-8") as handle: return json.load(handle) except Exception as exc: logger.debug("No %s for %s (%s)", filename, repo, exc) diff --git a/tests/parser/remote/test_crawler_markdown.py b/tests/parser/remote/test_crawler_markdown.py index a29ea03b..ca69e183 100644 --- a/tests/parser/remote/test_crawler_markdown.py +++ b/tests/parser/remote/test_crawler_markdown.py @@ -48,7 +48,7 @@ def _patch_validate_url(monkeypatch): @pytest.fixture(autouse=True) def _patch_tldextract(monkeypatch): monkeypatch.setattr( - "application.parser.remote.crawler_markdown.tldextract.extract", + "application.parser.remote.crawler_markdown._extract", _fake_extract, ) diff --git a/tests/parser/test_tokenization.py b/tests/parser/test_tokenization.py index afa7bd0f..e5a47ac0 100644 --- a/tests/parser/test_tokenization.py +++ b/tests/parser/test_tokenization.py @@ -1,6 +1,9 @@ """Chunk sizes must be counted in the embedding model's units, and splitting must never rewrite the text it splits.""" +import sys +import types + import pytest from application.parser import tokenization @@ -298,3 +301,35 @@ class TestUnknownTokenCollapse: assert "".join(pieces) == text, "split must not lose or alter text" assert len(pieces) > 1, "a collapsed run must still be cut into pieces" assert all(counter.count(p) <= 20 for p in pieces) + + +class TestTokenizerFile: + """The chunker's tokenizer must come from the hub cache without a network round trip.""" + + def test_cache_hit_makes_no_online_call(self, monkeypatch): + calls = [] + + def fake_download(repo, filename, local_files_only=False): + calls.append(local_files_only) + return "/cache/tokenizer.json" + + fake_hub = types.ModuleType("huggingface_hub") + fake_hub.hf_hub_download = fake_download + monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub) + assert tokenization._tokenizer_file("org/model") == "/cache/tokenizer.json" + assert calls == [True] + + def test_cache_miss_falls_back_to_online(self, monkeypatch): + calls = [] + + def fake_download(repo, filename, local_files_only=False): + calls.append(local_files_only) + if local_files_only: + raise FileNotFoundError("not cached") + return "/downloaded/tokenizer.json" + + fake_hub = types.ModuleType("huggingface_hub") + fake_hub.hf_hub_download = fake_download + monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub) + assert tokenization._tokenizer_file("org/model") == "/downloaded/tokenizer.json" + assert calls == [True, False] diff --git a/tests/scripts/test_prefetch_models.py b/tests/scripts/test_prefetch_models.py index ea8f5708..7892e0a8 100644 --- a/tests/scripts/test_prefetch_models.py +++ b/tests/scripts/test_prefetch_models.py @@ -94,3 +94,18 @@ class TestMain: with patch.object(prefetch_models, "prefetch", return_value=[]) as spy: prefetch_models.main(["granite-97m"]) assert spy.call_args.args[1] == "/app/models" + + +class TestPrefetchTiktoken: + def test_warms_every_listed_encoding(self): + """The image sets TIKTOKEN_CACHE_DIR; warming fills it at build time.""" + fake = MagicMock() + module = types.ModuleType("tiktoken") + module.get_encoding = fake + with patch.dict(sys.modules, {"tiktoken": module}): + fetched = prefetch_models.prefetch_tiktoken() + assert fetched == list(prefetch_models.TIKTOKEN_ENCODINGS) + assert [c.args[0] for c in fake.call_args_list] == list(prefetch_models.TIKTOKEN_ENCODINGS) + + def test_cl100k_is_the_encoding_token_counting_uses(self): + assert "cl100k_base" in prefetch_models.TIKTOKEN_ENCODINGS