"""Text-embedding cache and transient-encoder contract for Vili's resident Kimodo worker. The worker never holds the 8B text encoder. CachedTextEncoder answers Kimodo's encoder call from a CultCache store of float32 vectors and sends only the missing texts to kimodo-text-encode.py, a one-shot child that loads the encoder, encodes, writes its answer to a file and exits. Nothing here logs or raises with a prompt in it: an error names a count, a size or a rule, never a text. """ import hashlib import json import subprocess import sys import tempfile import time from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path import msgpack import torch from cultcache_py import CultCache, SingleFileMessagePackBackingStore, define_database_entry_type, database_entry_field, leaves as t GIB = 1024 ** 3 EMBEDDING_DIM = 4096 EMBEDDING_BYTES = EMBEDDING_DIM * 4 CHILD_PATH = Path(__file__).with_name("kimodo-text-encode.py") LLAMA_BASE = "meta-llama/Meta-Llama-3-8B-Instruct" # the base model both LLM2Vec adapters are applied to ASSEMBLY = "llm2vec-B-mntp-merged-then-supervised" MAX_THREADS = 16 def thread_cap(environ) -> int: """VILI_THREADS, the one torch thread budget of the worker and its encoder child: an integer from 1 to MAX_THREADS.""" raw = environ.get("VILI_THREADS", "") if not raw.isascii() or not raw.isdigit() or not 1 <= int(raw) <= MAX_THREADS: raise ValueError(f"VILI_THREADS must be an integer from 1 to {MAX_THREADS}") return int(raw) def apply_thread_cap(environ) -> int: threads = thread_cap(environ) torch.set_num_threads(threads) return threads def read_snapshot_revision(hf_home: Path, repo_id: str) -> str: """The commit the mounted Hugging Face cache resolves a repo's main ref to.""" ref = Path(hf_home) / "hub" / ("models--" + repo_id.replace("/", "--")) / "refs" / "main" return ref.read_text().strip() def snapshot_path(hf_home: Path, repo_id: str) -> Path: """The directory of the snapshot read_snapshot_revision names: assembly B loads from these, never from repo ids.""" return Path(hf_home) / "hub" / ("models--" + repo_id.replace("/", "--")) / "snapshots" / read_snapshot_revision(hf_home, repo_id) def encoder_identity(*, assembly: str, device: str, revisions: dict, dtype: str, torch_version: str, transformers_version: str, peft_version: str, accelerate_version: str, kimodo_commit: str, loader_sha256: str) -> str: """The canonical string a vector is keyed under: every input that can change it, in a fixed field order. Thread count is provenance, not identity; the encoder child's own bytes are, through loader_sha256.""" parts = [f"{repo}@{revision}" for repo, revision in sorted(revisions.items())] parts += [ f"assembly={assembly}", f"device={device}", f"dtype={dtype}", f"torch={torch_version}", f"transformers={transformers_version}", f"peft={peft_version}", f"accelerate={accelerate_version}", f"kimodo={kimodo_commit}", f"loader={loader_sha256}", ] return ";".join(parts) def embedding_key(identity: str, text: str) -> str: return hashlib.sha256((identity + "\n" + text).encode("utf-8")).hexdigest() @dataclass class TextEmbedding: encoderIdentity: str text: str dtype: str dim: int vector: bytes vectorSha256: str gpuBytes: int cpuBytes: int encodeSeconds: float createdAt: str TEXT_EMBEDDING = define_database_entry_type( "gamecult.vili.text_embedding.v0", [ database_entry_field("encoderIdentity", 0, codec=t.string), database_entry_field("text", 1, codec=t.string), database_entry_field("dtype", 2, codec=t.string), database_entry_field("dim", 3, codec=t.int32), database_entry_field("vector", 4, codec=t.bytes), database_entry_field("vectorSha256", 5, codec=t.string), database_entry_field("gpuBytes", 6, codec=t.int64), database_entry_field("cpuBytes", 7, codec=t.int64), database_entry_field("encodeSeconds", 8, codec=t.float64), database_entry_field("createdAt", 9, codec=t.string), ], cls=TextEmbedding, ) def open_store(path) -> CultCache: cache = ( CultCache.builder() .register_document_type(TEXT_EMBEDDING) .add_generic_store(SingleFileMessagePackBackingStore(path)) .build() ) cache.pull_all_backing_stores() return cache def run_child(texts: list, out_path, child_path=CHILD_PATH, timeout: float = 1800.0) -> dict: """Run the one-shot encoder child for `texts` and return its answer file. Raises RuntimeError on any non-zero exit.""" process = subprocess.run( [sys.executable, str(child_path), "--out", str(out_path)], input=json.dumps({"texts": texts}).encode("utf-8"), stdout=subprocess.DEVNULL, timeout=timeout, ) answer = msgpack.unpackb(Path(out_path).read_bytes(), raw=False) if Path(out_path).exists() else {} if process.returncode != 0: raise RuntimeError(f"text encoder child failed with exit code {process.returncode}") return answer def encode_misses_in_child(texts: list) -> dict: with tempfile.TemporaryDirectory() as directory: return run_child(texts, Path(directory) / "answer.mp") class CachedTextEncoder: """Kimodo's encoder contract, answered from the store; misses go to one child call per request.""" def __init__(self, store: CultCache, identity: str, encode_misses, device="cpu", dtype_name="float32") -> None: self.store = store self.identity = identity self._encode_misses = encode_misses self._device = device self._dtype_name = dtype_name self.reset() def reset(self) -> None: self.hits = 0 self.misses = 0 self.encode_seconds = 0.0 self.embeddings = [] def to(self, device): self._device = device return self def eval(self): return self def get_device(self): return self._device def count(self) -> int: return len(self.store.get_all(TEXT_EMBEDDING)) def __call__(self, texts): single = isinstance(texts, str) texts = [texts] if single else list(texts) keys = [embedding_key(self.identity, text) for text in texts] present = [self.store.get(TEXT_EMBEDDING, key) is not None for key in keys] missing = list(dict.fromkeys(text for text, found in zip(texts, present) if not found)) self.hits += sum(present) self.misses += len(missing) if missing: started = time.perf_counter() answer = self._encode_misses(missing) self.encode_seconds += time.perf_counter() - started self._admit(missing, answer) rows = [] for text, key in zip(texts, keys): record = self.store.get_required(TEXT_EMBEDDING, key) self.embeddings.append({"text": text, "vectorSha256": record.vectorSha256}) rows.append(torch.frombuffer(bytearray(record.vector), dtype=torch.float32)) features = torch.stack(rows)[:, None].to(self._device) return features, [1] * len(texts) def _admit(self, missing, answer) -> None: vectors = answer.get("vectors", []) if len(vectors) != len(missing) or any(len(v) != EMBEDDING_BYTES for v in vectors): raise RuntimeError(f"text encoder child returned {len(vectors)} vectors for {len(missing)} texts, each {EMBEDDING_BYTES} bytes expected") placement = answer.get("placement", {}) seconds = answer.get("encodeSeconds", [0.0] * len(missing)) for text, vector, took in zip(missing, vectors, seconds): vector = bytes(vector) self.store.put(TEXT_EMBEDDING, embedding_key(self.identity, text), TextEmbedding( encoderIdentity=self.identity, text=text, dtype="float32", dim=EMBEDDING_DIM, vector=vector, vectorSha256=hashlib.sha256(vector).hexdigest(), gpuBytes=0, cpuBytes=int(placement.get("cpuBytes", 0)), encodeSeconds=float(took), createdAt=datetime.now(timezone.utc).isoformat(), ))