mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 08:17:34 +00:00
fix(parser,vectorstore): stop the first-request downloads in a warmed install
Three things still reached the network from a container whose models were baked in: - tiktoken fetched cl100k_base from openaipublic.blob.core.windows.net on every fresh container (its cache defaulted to /tmp), and token accounting calls it on every chat. prefetch_models now warms it too; the image sets TIKTOKEN_CACHE_DIR. - The chunker loaded its tokenizer with Tokenizer.from_pretrained, which revalidates the revision with a HEAD request per process start and stalls for the etag timeout (10 s) when huggingface.co is unreachable. It now reads tokenizer.json from the hub cache first and only downloads on a miss; the repo-metadata read for models outside the registry does the same. - tldextract fetched the public suffix list on the first web crawl; the bundled snapshot is used instead. application/scripts/verify_offline.py exercises these paths (and docling's conversion when the extra is installed) so an image can be checked with docker run --network none.
This commit is contained in:
1 parent
196865846c
commit
4707c45b93
8 files changed
+262
-10
No files matched your search
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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<</Type/Catalog/Pages 2 0 R>>endobj\n"
|
||||
b"2 0 obj<</Type/Pages/Kids[3 0 R]/Count 1>>endobj\n"
|
||||
b"3 0 obj<</Type/Page/Parent 2 0 R/MediaBox[0 0 300 144]/Contents 4 0 R"
|
||||
b"/Resources<</Font<</F1 5 0 R>>>>>>endobj\n"
|
||||
b"4 0 obj<</Length 58>>stream\n"
|
||||
b"BT /F1 18 Tf 20 100 Td (Offline verification page) Tj ET\n"
|
||||
b"endstream\nendobj\n"
|
||||
b"5 0 obj<</Type/Font/Subtype/Type1/BaseFont/Helvetica>>endobj\n"
|
||||
b"trailer<</Root 1 0 R>>\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:]))
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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]
|
||||
@@ -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
|
||||
Reference in new issue
Block a user