docker/hf-space/check_image_cache.py

"""The vision checks' prepared asset views, kept on disk for every deepsearch-api process of the Space.

Before each check, usdsearch's llm_client prepares the asset's stored view: it decodes the stored thumbnail, fits it
to USDSEARCH_VISION_VALIDATION_IMAGE_MAX_PX and re-encodes it as base64 PNG, or JPEG with
USDSEARCH_VISION_VALIDATION_IMAGE_FORMAT=jpeg as on the Space (``Validation._encode_reference_images`` ->
``_encode_image``). That is most of a check's API CPU, and a stored view never changes. With
``CHECK_IMAGE_CACHE_DIR`` set (a Space Variable, pipeline/config.py), each view ``_encode_reference_images`` prepares
is kept in ``<dir>/<settings>/<sha256 of the stored base64>.b64``: a hit reads the file; a miss prepares the view as
before and writes it (a temp file, then ``os.replace``), so the API processes share one copy. ``<settings>`` is a
digest of everything else that can shape a view: the validation config, llm_client's validation module, Pillow and
its codecs. Query images are never cached. Each file carries its length and CRC-32. Any error (unreadable,
unwritable, full disk, a damaged file) falls back to preparing the view, so a check never fails here, and a damaged
file is written again. The views stay under MAX_BYTES and never leave less than MIN_FREE on the disk, which the
Space's index, media and caches share: past either, new views are prepared per check, as without the cache.

Unset or empty, nothing is imported or installed and checks run exactly as usdsearch's own code.
check_image_cache.pth, in the venv's site-packages, imports this module in every /app/.venv python (PYTHONPATH is
unset for them) while the variable is set. The import only arms a hook that wraps llm_client's Validation when, and
if, the process imports it. A Validation without the wrapped methods (a future wheel) logs one warning and runs
unchanged.

``python -m check_image_cache`` fills the cache once the index is in (supervisord-gpu.conf's
[program:check-image-cache]): every image of the image cache index, loaded as a check loads it and prepared through
the same wrapped path, so its files are the ones a check would write. It first removes what an earlier start left
(other settings' directories, stale temp files). ``--self-test`` is the image build's check (Dockerfile.hf-space-gpu).
"""

import functools
import hashlib
import json
import logging
import os
import re
import shutil
import sys
import threading
import time
import weakref
import zlib

ENV = "CHECK_IMAGE_CACHE_DIR"
TARGET = "llm_client.validation.validation"
MAX_BYTES = 2_000_000_000  # 448-px views: ~28 KB as JPEG, the Space's ~2,600 images ~75 MB (PNG ~121 KB, ~340 MB)
MIN_FREE = 12_000_000_000  # above the 8 GB the media loader keeps free
RECHECK_BYTES, RECHECK_S = 64_000_000, 60  # each process re-reads the size and the free space this often
_DIGEST = re.compile(r"^[0-9a-f]{24}$")
_TMP = re.compile(r"^[0-9a-f]{64}\.b64\.\d+\.\d+\.tmp$")
_MARK = "__check_image_cache__"
logger = logging.getLogger("check_image_cache")
_warned = set()


def _warn_once(kind, message, *args):
    if kind not in _warned:
        _warned.add(kind)
        logger.warning(message, *args)


def settings_digest(validation, module_file):
    """What, besides the stored image, shapes a prepared view: one directory per combination."""
    try:
        from PIL import __version__ as pillow, features
    except ImportError:  # no Pillow, no views to prepare (the tests' stand-in Validation)
        pillow = features = None
    config = validation._config
    dump = config.model_dump() if hasattr(config, "model_dump") else dict(vars(config))
    codecs = {}
    for name in ("jpg", "zlib", "webp"):
        try:
            codecs[name] = features.version(name)
        except Exception:
            codecs[name] = None
    with open(module_file, "rb") as fh:
        module = hashlib.sha256(fh.read()).hexdigest()
    identity = {
        "config": dump,  # all of it: a setting llm_client reads later can never share a stale directory
        "class": f"{type(validation).__module__}.{type(validation).__qualname__}",
        "module": module,
        "pillow": pillow,
        "codecs": codecs,
        "python": "%d.%d" % sys.version_info[:2],
    }
    return hashlib.sha256(json.dumps(identity, sort_keys=True, default=str).encode()).hexdigest()[:24]


class Cache:
    """One Validation's files, ``<root>/<settings digest>/<key>.b64``. Never raises: a failure is a miss."""

    def __init__(self, directory):
        self.directory = directory
        self._lock = threading.Lock()
        self.counts = {"hit": 0, "written": 0, "not_kept": 0, "bypass": 0}
        self._used = self._free = None  # bytes of views and free on the disk, as last read plus this process's writes
        self._since, self._read_at = 0, 0.0

    def _count(self, outcome):
        with self._lock:
            self.counts[outcome] += 1

    def _room(self, size):
        """True while the views stay under MAX_BYTES and the disk keeps MIN_FREE: both re-read from the disk every
        RECHECK_BYTES this process writes or RECHECK_S, so all processes together overshoot by little."""
        with self._lock:
            now = time.monotonic()
            if self._used is None or self._since >= RECHECK_BYTES or now - self._read_at >= RECHECK_S:
                self._used, self._free = files(self.directory)[1], _free(self.directory)
                self._since, self._read_at = 0, now
            if self._used + size > MAX_BYTES or self._free - size < MIN_FREE:
                return False
            self._used += size
            self._free -= size
            self._since += size
            return True

    def view(self, validation, image):
        """``validation._encode_image(image)``, read from its file when there is one."""
        path = self._path(validation, image)
        if path is None:
            self._count("bypass")
            return validation._encode_image(image)
        value = self._read(path)
        if value is not None:
            self._count("hit")
            return value
        value = validation._encode_image(image)  # its errors are the check's, as without the cache
        self._count("written" if self._write(path, value) else "not_kept")
        return value

    def _path(self, validation, image):
        # Only stored base64 images: a str that is not one is a file path, whose content can change under its name.
        try:
            if isinstance(image, str) and validation._is_base64(image):
                return os.path.join(self.directory, hashlib.sha256(image.encode()).hexdigest() + ".b64")
        except Exception as e:
            _warn_once("key", "check image cache: no key for an image (%s); preparing it per check", e)
        return None

    def _read(self, path):
        try:
            with open(path, "rb") as fh:
                data = fh.read()
        except FileNotFoundError:
            return None
        except Exception as e:
            _warn_once("read", "check image cache: cannot read %s (%s); preparing views per check", path, e)
            return None
        head, _, body = data.partition(b"\n")
        try:
            size, crc = head.split(b" ")
            if int(size) == len(body) and int(crc, 16) == zlib.crc32(body):
                return body.decode("ascii")
        except (ValueError, UnicodeDecodeError):
            pass
        _warn_once("damaged", "check image cache: %s is damaged; preparing the view again", path)
        return None

    def _write(self, path, value):
        tmp = f"{path}.{os.getpid()}.{threading.get_ident()}.tmp"
        try:
            body = value.encode("ascii")
            if not self._room(len(body) + 32):
                _warn_once("full", "check image cache: full (%d MB of views, or under %d GB free on the disk); "
                           "new views are prepared per check", MAX_BYTES // 1_000_000, MIN_FREE // 1_000_000_000)
                return False
            try:
                fh = open(tmp, "wb")
            except FileNotFoundError:
                os.makedirs(self.directory, exist_ok=True)
                fh = open(tmp, "wb")
            with fh:
                fh.write(b"%d %08x\n" % (len(body), zlib.crc32(body)))
                fh.write(body)
            os.replace(tmp, path)
            return True
        except Exception as e:
            _warn_once("write", "check image cache: cannot write in %s (%s); views are prepared per check",
                        self.directory, e)
            try:
                os.unlink(tmp)
            except OSError:
                pass
            return False


_caches = weakref.WeakKeyDictionary()
_caches_lock = threading.Lock()


def is_wrapped(cls):
    return getattr(getattr(cls, "_encode_reference_images", None), _MARK, None) is not None


def cache_of(validation):
    """The Cache of a Validation whose class is wrapped, else None (also when its settings can't be read)."""
    mark = getattr(getattr(type(validation), "_encode_reference_images", None), _MARK, None)
    if mark is None:
        return None
    try:
        return _caches[validation]
    except KeyError:
        pass
    except TypeError:  # not weak-referenceable: no cache
        return None
    root, module_file = mark
    try:
        cache = Cache(os.path.join(root, settings_digest(validation, module_file)))
    except Exception as e:
        _warn_once("settings", "check image cache off: cannot read the view settings (%s)", e)
        cache = None
    with _caches_lock:
        return _caches.setdefault(validation, cache)


def patch(module, root):
    """Wrap ``module.Validation._encode_reference_images`` (llm_client's validation module); True when wrapped."""
    cls = getattr(module, "Validation", None)
    original = getattr(cls, "_encode_reference_images", None)
    if original is None or not all(callable(getattr(cls, n, None)) for n in ("_encode_image", "_is_base64")):
        _warn_once("missing", "check image cache off: %s.Validation has no _encode_reference_images, "
                   "_encode_image or _is_base64", getattr(module, "__name__", module))
        return False
    if getattr(original, _MARK, None) is not None:
        return True

    @functools.wraps(original)
    def _encode_reference_images(self, reference_images):
        cache = cache_of(self) if reference_images else None
        if cache is None:
            return original(self, reference_images)
        return [cache.view(self, image) for image in reference_images]

    setattr(_encode_reference_images, _MARK, (root, getattr(module, "__file__", None)))
    cls._encode_reference_images = _encode_reference_images
    return True


class _PatchOnImport:
    """A sys.meta_path finder: wraps llm_client's Validation right after its module has run, in the process that
    imports it. Every other import passes through untouched."""

    def __init__(self, root):
        self.root = root

    def find_spec(self, name, path, target=None):
        if name != TARGET:
            return None
        try:
            from importlib.machinery import PathFinder, SourceFileLoader, SourcelessFileLoader

            spec = PathFinder.find_spec(name, path)
            if spec is None:
                return None
            # A file loader serves this one module, so its exec_module can be wrapped in place.
            if type(spec.loader) not in (SourceFileLoader, SourcelessFileLoader):
                _warn_once("hook", "check image cache off: %s is not loaded from a file", name)
                return None
            exec_module = spec.loader.exec_module
            root = self.root

            def exec_and_patch(module):
                exec_module(module)
                try:
                    patch(module, root)
                except Exception as e:
                    _warn_once("patch", "check image cache off: %s", e)

            spec.loader.exec_module = exec_and_patch
            return spec
        except Exception as e:
            _warn_once("hook", "check image cache off: %s", e)
            return None


def install(environ=os.environ):
    """Arm the hook when CHECK_IMAGE_CACHE_DIR is set; with it unset or empty, nothing changes."""
    root = (environ.get(ENV) or "").strip()
    if not root:
        return False
    module = sys.modules.get(TARGET)
    if module is not None:
        return patch(module, root)
    if not any(isinstance(f, _PatchOnImport) for f in sys.meta_path):
        sys.meta_path.insert(0, _PatchOnImport(root))
    return True


async def fill(validation, loader, threads=4, page=500, batch=25):
    """Prepare every image of the loader's index through ``validation._encode_reference_images``.

    The ids come from a scroll of the index, the images from ``loader.load_images`` (what a check loads), and up to
    ``threads`` views are prepared at once in threads. Returns {"images", "empty", "failed"}.
    """
    import asyncio
    from concurrent.futures import ThreadPoolExecutor

    counts = {"images": 0, "empty": 0, "failed": 0}
    loop = asyncio.get_running_loop()
    slots = asyncio.Semaphore(threads * 2)
    pending = set()

    def done(task):  # on the loop's thread
        slots.release()
        pending.discard(task)
        if task.cancelled() or task.exception() is not None:
            counts["failed"] += 1

    with ThreadPoolExecutor(threads, thread_name_prefix="check-image-cache") as pool:
        client, index = loader.client, loader.image_cache_index
        resp = await client.search(index=index, body={"query": {"match_all": {}}, "_source": False, "sort": ["_doc"]},
                                   size=page, scroll="5m")
        try:
            while resp["hits"]["hits"]:
                ids = [hit["_id"] for hit in resp["hits"]["hits"]]
                for i in range(0, len(ids), batch):
                    images = await loader.load_images(ids[i:i + batch])
                    for image_id in ids[i:i + batch]:
                        image = images.get(image_id)
                        if not isinstance(image, str) or not image:
                            counts["empty"] += 1
                            continue
                        counts["images"] += 1
                        await slots.acquire()
                        task = loop.run_in_executor(pool, validation._encode_reference_images, [image])
                        pending.add(task)
                        task.add_done_callback(done)
                resp = await client.scroll(scroll_id=resp["_scroll_id"], scroll="5m")
        finally:
            if pending:
                await asyncio.gather(*pending, return_exceptions=True)
            try:
                await client.clear_scroll(scroll_id=resp["_scroll_id"])
            except Exception:
                pass
    return counts


def files(directory):
    """(count, bytes) of the views in a settings directory."""
    n = size = 0
    try:
        with os.scandir(directory) as it:
            for entry in it:
                if entry.name.endswith(".b64"):
                    try:
                        size += entry.stat().st_size
                        n += 1
                    except OSError:  # replaced or removed meanwhile
                        pass
    except OSError:
        pass
    return n, size


def _free(path):
    """Free bytes on the disk that holds ``path`` (its nearest existing parent); 0 when unreadable: no writes."""
    try:
        while not os.path.exists(path) and os.path.dirname(path) != path:
            path = os.path.dirname(path)
        return shutil.disk_usage(path).free
    except OSError:
        return 0


def sweep(directory, older_than=600):
    """What an earlier start left in the cache's root: other settings' directories, and temp files of writers that
    died (older than ``older_than`` s). Only names this module makes are touched. Returns the number removed."""
    root, removed = os.path.dirname(directory), 0
    try:
        with os.scandir(root) as it:
            for entry in it:
                if _DIGEST.match(entry.name) and entry.path != directory and entry.is_dir(follow_symlinks=False):
                    shutil.rmtree(entry.path, ignore_errors=True)
                    removed += 1
        with os.scandir(directory) as it:
            for entry in it:
                if _TMP.match(entry.name) and time.time() - entry.stat().st_mtime > older_than:
                    os.unlink(entry.path)
                    removed += 1
    except OSError:
        pass
    return removed


def self_test():
    """The image build's check (Dockerfile.hf-space-gpu): the .pth loaded this module and it wrapped the vendored
    llm_client, views read back from the cache are byte-identical to llm_client's own and in the format the
    environment names (USDSEARCH_VISION_VALIDATION_IMAGE_FORMAT, unset: PNG), and the warm-up's imports exist.
    Raises on any difference. Writes only in CHECK_IMAGE_CACHE_DIR, which the build then removes."""
    import base64
    import io

    from PIL import Image
    from deepsearch_api.routers_v2.dependencies import build_image_loader  # noqa: F401 (the warm-up's imports)
    from deepsearch_api.search_backend.main import SearchSettings  # noqa: F401
    from deepsearch_api.validation.config import ValidationSettings
    from llm_client.validation.validation import Validation

    assert "check_image_cache" in sys.modules and is_wrapped(Validation), "the hook did not wrap llm_client"
    global MIN_FREE
    MIN_FREE = 0  # the image builder's disk is not the Space's

    def b64(image, fmt, **kw):
        buf = io.BytesIO()
        image.save(buf, fmt, **kw)
        return base64.b64encode(buf.getvalue()).decode()

    noise = Image.effect_noise((1024, 1024), 40)
    rgb = Image.merge("RGB", (noise, Image.linear_gradient("L").resize((1024, 1024)), noise.rotate(90)))
    rgba = rgb.resize((900, 600))
    rgba.putalpha(Image.linear_gradient("L").resize((900, 600)))
    images = [b64(rgb, "JPEG", quality=90), b64(rgba, "PNG"), b64(rgb.resize((300, 200)), "JPEG", quality=90)]
    config = ValidationSettings(image_max_px=448).to_validation_config()
    plain = [Validation(config=config)._encode_image(image) for image in images]
    # Wheels that dropped the format setting would pass the build's JPEG run with PNG views.
    wanted = os.environ.get("USDSEARCH_VISION_VALIDATION_IMAGE_FORMAT") or "png"
    magic = {"png": "iVBORw0KGgo", "jpeg": "/9j/"}[wanted]
    assert all(view.startswith(magic) for view in plain), f"the views are not {wanted}"
    miss = Validation(config=config)._encode_reference_images(images)
    reader = Validation(config=config)
    hit = reader._encode_reference_images(images)
    assert miss == plain and hit == plain, "a cached view differs from llm_client's own"
    assert cache_of(reader).counts["hit"] == len(images), "the views were not read back from the cache"
    return len(images)


def main(argv=None):
    """``python -m check_image_cache [--threads N]``: fill the cache from the image cache index. Always exits 0
    (``--self-test`` exits 1 when the check fails)."""
    import argparse
    import asyncio
    import resource

    parser = argparse.ArgumentParser(prog="check_image_cache")
    parser.add_argument("--threads", type=int, default=4)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args(argv)
    say = functools.partial(print, "check-image-cache:", flush=True)
    if not (os.environ.get(ENV) or "").strip():
        say(f"{ENV} is not set: nothing to do")
        return 1 if args.self_test else 0
    if args.self_test:
        say(f"self-test passed: {self_test()} views read back byte-identical")
        return 0
    started, cpu = time.monotonic(), time.process_time()
    try:
        from deepsearch_api.validation.config import ValidationSettings

        settings = ValidationSettings()
        if not settings.enabled:
            say("vision validation is off: nothing to do")
            return 0
        from llm_client.validation.validation import Validation

        validation = Validation(config=settings.to_validation_config())
        cache = cache_of(validation)
        if cache is None:
            say("llm_client's Validation is not wrapped (see the warning above): nothing to do")
            return 0
        swept = sweep(cache.directory)
        from deepsearch_api.routers_v2.dependencies import build_image_loader
        from deepsearch_api.search_backend.main import SearchSettings

        loader = build_image_loader(SearchSettings())

        async def run():
            try:
                return await fill(validation, loader, threads=max(1, args.threads))
            finally:
                await loader.client.close()

        counts = asyncio.run(run())
    except Exception as e:
        say(f"stopped ({type(e).__name__}: {e}); checks prepare the rest as they go")
        return 0
    n, size = files(cache.directory)
    c = cache.counts
    say(f"{counts['images']} images in {time.monotonic() - started:.1f} s, {time.process_time() - cpu:.1f} CPU-s: "
        f"{c['written']} written, {c['hit']} already there, {c['not_kept']} not kept, {c['bypass']} not cacheable, "
        f"{counts['failed']} failed, {counts['empty']} missing or not loaded; {n} files, {size / 1e6:.0f} MB in "
        f"{cache.directory} ({swept} old entries removed); "
        f"peak RSS {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024:.0f} MB")
    return 0


if __name__ == "__main__":
    import check_image_cache  # the instance the .pth imported: its hook wrapped Validation

    sys.exit(check_image_cache.main())
else:
    try:
        install()
    except Exception as e:  # never break a python start
        _warn_once("install", "check image cache off: %s", e)