Spaces:
Running on RTX PRO 6000
Running on RTX PRO 6000
| """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 | |
| 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) | |