dp-zp-agent/tools/runtime/runtime_assets.py

"""Prepare immutable model/tokenizer assets during build, never during startup."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

DEFAULT_MODEL = "intfloat/multilingual-e5-small"
DEFAULT_REVISION = "614241f622f53c4eeff9890bdc4f31cfecc418b3"
DEFAULT_ROOT = Path(__file__).resolve().parents[2] / "runtime_assets"


def asset_root() -> Path:
    return Path(os.getenv("ZP_RUNTIME_ASSETS", str(DEFAULT_ROOT)))


def manifest(root: Path | None = None) -> dict | None:
    path = (root or asset_root()) / "manifest.json"
    if not path.exists():
        if os.getenv("ZP_ASSETS_REQUIRED") == "1":
            raise RuntimeError(
                "Chýbajú runtime assets; zostavte image alebo spustite python -m tools.runtime.runtime_assets prepare."
            )
        return None
    return json.loads(path.read_text())


def configure_tokenizer() -> None:
    data = manifest()
    if data is not None:
        # Bundled immutable files take precedence over an empty user cache.
        os.environ["TIKTOKEN_CACHE_DIR"] = str(asset_root() / "tiktoken")


def local_model(model_name: str, revision: str | None) -> str | None:
    data = manifest()
    if data and data["embedding"]["model"] == model_name:
        if revision and revision != data["embedding"]["revision"]:
            raise RuntimeError(
                "Požadovaná revízia modelu nie je v image; vytvorte nový image s touto revíziou."
            )
        return str(asset_root() / "model")
    if data and os.getenv("ZP_ASSETS_REQUIRED") == "1":
        raise RuntimeError("Požadovaný embeddingový model nie je v image; vytvorte nový image.")
    return None


def verify(root: Path) -> dict:
    data = json.loads((root / "manifest.json").read_text())
    for relative, digest in data["files"].items():
        path = root / relative
        if not path.resolve().is_relative_to(root.resolve()):
            raise RuntimeError("Manifest obsahuje cestu mimo runtime assets")
        if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest() != digest:
            raise RuntimeError(f"Chýbajúci alebo poškodený runtime asset: {relative}")
    if not any(name.startswith("tiktoken/") for name in data["files"]):
        raise RuntimeError("Manifest neobsahuje tokenizér")
    required = {"model/config.json", "model/modules.json", "model/tokenizer.json"}
    if not required.issubset(data["files"]) or not any(
        name.endswith(".safetensors") for name in data["files"]
    ):
        raise RuntimeError("Manifest neobsahuje kompletný embeddingový model")
    return data


def offline_probe(root: Path) -> None:
    """Fresh interpreter, with network calls forbidden even on an online builder."""
    environment = {
        **os.environ,
        "ZP_RUNTIME_ASSETS": str(root.resolve()),
        "ZP_ASSETS_REQUIRED": "1",
        "HF_HUB_OFFLINE": "1",
        "TRANSFORMERS_OFFLINE": "1",
        "HF_HOME": str(root / "unused-user-cache"),
        "TIKTOKEN_CACHE_DIR": str(root / "unused-token-cache"),
    }
    code = """
import socket
def denied(*args, **kwargs): raise RuntimeError('Network access during offline asset probe')
socket.socket.connect = denied
socket.socket.connect_ex = denied
socket.getaddrinfo = denied
from tools.runtime.runtime_assets import configure_tokenizer, manifest, local_model
configure_tokenizer()
import tiktoken
for encoding in manifest()['encodings']:
    assert tiktoken.get_encoding(encoding).encode('Overenie štartu bez siete.')
from sentence_transformers import SentenceTransformer
data = manifest()['embedding']
model = SentenceTransformer(local_model(data['model'], data['revision']), local_files_only=True)
assert model.encode(['query: overenie bez siete'], show_progress_bar=False).shape[0] == 1
print('Offline model/tokenizer probe: OK')
"""
    subprocess.run([sys.executable, "-c", code], env=environment, check=True)


def prepare(root: Path, model: str, revision: str | None, encodings: list[str]) -> dict:
    revision = revision or (DEFAULT_REVISION if model == DEFAULT_MODEL else "")
    if not re.fullmatch("[0-9a-f]{40}", revision):
        raise ValueError("Príprava modelu vyžaduje presnú 40-znakovú Git revíziu")
    if (root / "manifest.json").is_file():
        data = verify(root)
        if (
            data["embedding"] == {"model": model, "revision": revision}
            and data["encodings"] == encodings
        ):
            offline_probe(root)
            return data
        raise RuntimeError(
            "Assets patria inej konfigurácii; použite nový/prázdny pracovný adresár."
        )
    from huggingface_hub import snapshot_download
    from huggingface_hub.errors import LocalEntryNotFoundError

    try:
        source = snapshot_download(model, revision=revision, local_files_only=True)
    except (FileNotFoundError, LocalEntryNotFoundError):
        source = snapshot_download(
            model,
            revision=revision,
            allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt"],
            ignore_patterns=["onnx/*", "openvino/*"],
        )
    root.mkdir(parents=True, exist_ok=True)
    shutil.copytree(
        source,
        root / "model",
        dirs_exist_ok=True,
        symlinks=False,
        ignore=shutil.ignore_patterns(
            "onnx", "openvino", "*.bin", "*.md", ".cache", ".gitattributes"
        ),
    )
    environment = {**os.environ, "TIKTOKEN_CACHE_DIR": str(root.resolve() / "tiktoken")}
    subprocess.run(
        [
            sys.executable,
            "-c",
            "import sys,tiktoken; [tiktoken.get_encoding(name) for name in sys.argv[1:]]",
            *encodings,
        ],
        env=environment,
        check=True,
    )
    files = {
        str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
        for folder in ["model", "tiktoken"]
        for p in sorted((root / folder).rglob("*"))
        if p.is_file()
    }
    data = {
        "schema_version": 1,
        "embedding": {"model": model, "revision": revision},
        "encodings": encodings,
        "files": files,
    }
    (root / "manifest.json").write_text(json.dumps(data, sort_keys=True, indent=2) + "\n")
    verify(root)
    offline_probe(root)
    return data


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["prepare", "verify"])
    parser.add_argument("--output", type=Path, default=asset_root())
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--revision")
    parser.add_argument("--encoding", action="append")
    args = parser.parse_args()
    if args.command == "prepare":
        data = prepare(
            args.output,
            args.model,
            args.revision,
            list(dict.fromkeys(args.encoding or ["cl100k_base"])),
        )
    else:
        data = verify(args.output)
        offline_probe(args.output)
    print(
        json.dumps(
            {
                "model": data["embedding"],
                "encodings": data["encodings"],
                "files": len(data["files"]),
            }
        )
    )


if __name__ == "__main__":
    main()