From c3684b733b680b2d644866c4d2dc20b8681583be Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Fri, 23 Oct 2026 12:14:09 +0200 Subject: [PATCH] Encoder: assembly B on CPU with a complete identity and a thread cap The encoder child is CPU only (GPU placement, GpuBusy and plan_placement are deleted) and assembles LLM2Vec B: Llama base, MNTP merged, then the supervised adapter. The cache identity names assembly, device, dtype, torch/transformers/peft/accelerate versions, Kimodo commit and the child's own sha256 beside the three revisions. VILI_THREADS (1..16) caps torch in the worker and the child. The worker binds port 0 and reports it in its ready line; errors name a class and place, never the request. Co-Authored-By: Claude Opus 5.5 --- .gitmodules | 3 + scripts/kimodo-resident-worker.py | 43 +++++++++----- scripts/kimodo-text-encode.py | 84 ++++++++------------------ scripts/test_text_encoding.py | 97 +++++++++++++++---------------- scripts/text_encoding.py | 49 ++++++++-------- vendor/CultLib | 1 + 6 files changed, 126 insertions(+), 151 deletions(-) create mode 100644 .gitmodules create mode 160000 vendor/CultLib diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..5098433 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "vendor/CultLib"] + path = vendor/CultLib + url = https://github.com/GameCult/CultLib.git diff --git a/scripts/kimodo-resident-worker.py b/scripts/kimodo-resident-worker.py index 846ace9..b479a03 100644 --- a/scripts/kimodo-resident-worker.py +++ b/scripts/kimodo-resident-worker.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import argparse +import hashlib import json import os import shutil @@ -9,6 +10,7 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from threading import Lock from types import SimpleNamespace +import accelerate import peft import torch import transformers @@ -34,32 +36,39 @@ import text_encoding def parse_args(): parser = argparse.ArgumentParser(description="Resident Vili worker for Kimodo motion generation") - parser.add_argument("--host", default="0.0.0.0") - parser.add_argument("--port", type=int, default=8825) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=0) parser.add_argument("--model", default=DEFAULT_MODEL) return parser.parse_args() args = parse_args() -device = "cuda:0" if torch.cuda.is_available() else "cpu" +device = "cpu" +threads = text_encoding.apply_thread_cap(os.environ) +torch.set_num_interop_threads(1) started_at = time.time() model_load_started_at = time.time() preset = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"] kimodo_commit = os.environ["VILI_KIMODO_COMMIT"] -kimodo_image = os.environ["VILI_KIMODO_IMAGE"] +python_env = os.environ["VILI_PYTHON_ENV"] encoder_identity = text_encoding.encoder_identity( + assembly=text_encoding.ASSEMBLY, + device=device, revisions={ repo: text_encoding.read_snapshot_revision(HF_HOME, repo) for repo in (preset["base_model_name_or_path"], preset["peft_model_name_or_path"], text_encoding.LLAMA_BASE) }, dtype=preset["dtype"], + torch_version=torch.__version__, transformers_version=transformers.__version__, peft_version=peft.__version__, + accelerate_version=accelerate.__version__, kimodo_commit=kimodo_commit, + loader_sha256=hashlib.sha256(text_encoding.CHILD_PATH.read_bytes()).hexdigest(), ) # The 8B encoder never lives here: misses go to a one-shot child (kimodo-text-encode.py), hits come from the store. text_encoder = text_encoding.CachedTextEncoder( - text_encoding.open_store(os.path.join(os.environ.get("VILI_EMBEDDINGS_DIR", "/vili-embeddings"), "text-embeddings.cc")), + text_encoding.open_store(os.path.join(os.environ["VILI_EMBEDDINGS_DIR"], "text-embeddings.cc")), encoder_identity, text_encoding.encode_misses_in_child, device=device, @@ -72,7 +81,7 @@ model, resolved_model = load_model( text_encoder=text_encoder, ) model_revision = text_encoding.read_snapshot_revision(HF_HOME, MODEL_NAMES[resolved_model]) -worker_identity = {"kimodoCommit": kimodo_commit, "imageId": kimodo_image, "modelRevision": model_revision} +worker_identity = {"kimodoCommit": kimodo_commit, "pythonEnv": python_env, "modelRevision": model_revision, "threads": threads} model_load_seconds = time.time() - model_load_started_at model_info = get_model_info(resolved_model) display_model = model_info.display_name if model_info else resolved_model @@ -272,10 +281,12 @@ def generate(payload): requested_model = payload.get("model") if requested_model and requested_model != resolved_model and requested_model != display_model: - raise ValueError(f"Worker model is {resolved_model}; request asked for {requested_model}.") + raise ValueError("The request names a model this worker does not serve.") meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} - output_base = str(payload.get("output") or "/outputs/motion") + if not payload.get("output"): + raise ValueError("The request carries no output path.") + output_base = str(payload["output"]) diffusion_steps = max(1, int(payload.get("diffusionSteps") or payload.get("diffusion_steps") or meta.get("diffusion_steps") or 20)) num_samples = max(1, int(payload.get("numSamples") or payload.get("num_samples") or meta.get("num_samples") or 1)) num_transition_frames = int(payload.get("numTransitionFrames") or payload.get("num_transition_frames") or 5) @@ -394,22 +405,24 @@ class Handler(BaseHTTPRequestHandler): payload = json.loads(self.rfile.read(length).decode("utf-8") or "{}") write_json(self, 200, generate(payload)) except Exception as error: - write_json(self, 500, { - "ok": False, - "error": str(error), - "traceback": traceback.format_exc(), - }) + # An error names its class and where it was raised, never a message: messages can carry the request. + frame = traceback.extract_tb(error.__traceback__)[-1] + where = f"{os.path.basename(frame.filename)}:{frame.lineno}" + print(f"generate failed: {type(error).__name__} at {where}", flush=True) + write_json(self, 500, {"ok": False, "error": f"{type(error).__name__} at {where}"}) def log_message(self, fmt, *values): print(f"{self.address_string()} - {fmt % values}", flush=True) +server = ThreadingHTTPServer((args.host, args.port), Handler) print(json.dumps({ "event": "kimodo-worker-ready", "device": device, "model": resolved_model, "displayModel": display_model, "modelLoadSeconds": model_load_seconds, - "port": args.port, + "port": server.server_address[1], + "threads": threads, }), flush=True) -ThreadingHTTPServer((args.host, args.port), Handler).serve_forever() +server.serve_forever() diff --git a/scripts/kimodo-text-encode.py b/scripts/kimodo-text-encode.py index ce91f18..6bd9e2a 100644 --- a/scripts/kimodo-text-encode.py +++ b/scripts/kimodo-text-encode.py @@ -1,12 +1,11 @@ #!/usr/bin/env python3 -"""One-shot Kimodo text encoder: texts on stdin, vectors in --out, then exit so the OS reclaims RAM and VRAM. +"""One-shot Kimodo text encoder: texts on stdin, vectors in --out, then exit so the OS reclaims its RAM. -Reads {"texts": [...]} as JSON on stdin. Loads Kimodo's LLM2Vec preset in bfloat16 with the layers that do not -fit in free VRAM held in RAM, encodes each text alone as LLM2VecEncoder does, and writes a MessagePack answer. -There is no CPU fallback; --device cpu exists only for the reference probe. +Reads {"texts": [...]} as JSON on stdin. Loads LLM2Vec assembly B on CPU in bfloat16, encodes each text alone as +LLM2VecEncoder does, and writes a MessagePack answer. The torch thread budget is VILI_THREADS. """ -import argparse import json +import os import resource import sys import time @@ -16,100 +15,63 @@ import msgpack import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) -from text_encoding import CHILD_EXIT_GPU_BUSY, EMBEDDING_DIM, GpuBusy, plan_placement # noqa: E402 +from text_encoding import ASSEMBLY, EMBEDDING_DIM, LLAMA_BASE, apply_thread_cap # noqa: E402 from peft import PeftModel # noqa: E402 from transformers import AutoConfig, AutoTokenizer # noqa: E402 -from kimodo.model.llm2vec.llm2vec import LLM2Vec, batch_to_device # noqa: E402 +from kimodo.model.llm2vec.llm2vec import LLM2Vec # noqa: E402 from kimodo.model.load_model import TEXT_ENCODER_PRESETS # noqa: E402 -def load_encoder(base_id, peft_id, dtype, peft_kwargs, **model_kwargs): - """Assemble the encoder in the order and with the ids of LLM2Vec.from_pretrained. +def load_encoder(mntp_id, supervised_id, dtype): + """Assembly B: the Llama base with the MNTP adapter merged in, then the supervised adapter on top. - MUST TRACK Kimodo's loader: kimodo/model/llm2vec/llm2vec.py, LLM2Vec.from_pretrained, Kimodo commit 6bb5848 - (nv-tlabs/kimodo). It is mirrored, not called, because that method hands PeftModel.from_pretrained no kwargs, and - with layers offloaded to RAM peft then takes a path that raises TypeError in accelerate's get_balanced_memory - (peft 0.18.1, accelerate 1.13.0, transformers 5.1.0). Here every PeftModel.from_pretrained gets peft_kwargs - (device_map='sequential'). Operator ruling motion-pipeline:ruling:encoder-own-loader. The ids' llm2vec_config.json - is not a local path, so Kimodo's loader applies no encoder config; neither does this. + Kimodo's own LLM2VecEncoder loads the MNTP repo as the base and never applies the MNTP weights under the + pinned transformers (assembly A); this assembly is ours (motion-pipeline:ruling:encoder-assembly-b). The + identity carries the Kimodo commit and this file's sha256, so a changed loader misses the cache. """ - tokenizer = AutoTokenizer.from_pretrained(base_id) + tokenizer = AutoTokenizer.from_pretrained(LLAMA_BASE) tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" - model_class = LLM2Vec._get_model_class(AutoConfig.from_pretrained(base_id).__class__.__name__, enable_bidirectional=True) - model = model_class.from_pretrained(base_id, torch_dtype=dtype, **model_kwargs) - if hasattr(model, "peft_config"): - model = PeftModel.from_pretrained(model, base_id, **peft_kwargs).merge_and_unload() - model = PeftModel.from_pretrained(model, peft_id, **peft_kwargs) + model_class = LLM2Vec._get_model_class(AutoConfig.from_pretrained(LLAMA_BASE).__class__.__name__, enable_bidirectional=True) + model = model_class.from_pretrained(LLAMA_BASE, torch_dtype=dtype) + model = PeftModel.from_pretrained(model, mntp_id).merge_and_unload() + model = PeftModel.from_pretrained(model, supervised_id) return LLM2Vec(model=model, tokenizer=tokenizer) -def mem_available() -> int: - for line in Path("/proc/meminfo").read_text().splitlines(): - if line.startswith("MemAvailable:"): - return int(line.split()[1]) * 1024 - raise RuntimeError("MemAvailable is not in /proc/meminfo") - - def main() -> int: - parser = argparse.ArgumentParser() - parser.add_argument("--out", required=True) - parser.add_argument("--device", choices=["gpu", "cpu"], default="gpu") - args = parser.parse_args() + apply_thread_cap(os.environ) texts = json.load(sys.stdin)["texts"] + out = Path(sys.argv[sys.argv.index("--out") + 1]) preset = TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"] - load_kwargs, peft_kwargs = {}, {} - if args.device == "gpu": - free_vram, _ = torch.cuda.mem_get_info(0) - try: - max_memory = plan_placement(free_vram, mem_available()) - load_kwargs = {"device_map": "auto", "max_memory": max_memory} - peft_kwargs = {"device_map": "sequential", "max_memory": max_memory} - except GpuBusy as busy: - Path(args.out).write_bytes(msgpack.packb({"error": busy.code, "freeVramBytes": busy.free_vram}, use_bin_type=True)) - return CHILD_EXIT_GPU_BUSY - started = time.perf_counter() - model = load_encoder( - preset["base_model_name_or_path"], - preset["peft_model_name_or_path"], - getattr(torch, preset["dtype"]), - peft_kwargs, - **load_kwargs, - ) + model = load_encoder(preset["base_model_name_or_path"], preset["peft_model_name_or_path"], getattr(torch, preset["dtype"])) model.eval() for p in model.parameters(): p.requires_grad = False load_seconds = time.perf_counter() - started - placement = {"gpuBytes": 0, "cpuBytes": 0} - for p in model.parameters(): - placement["gpuBytes" if p.device.type == "cuda" else "cpuBytes"] += p.numel() * p.element_size() - input_device = "cuda:0" if args.device == "gpu" else "cpu" - vectors, encode_seconds = [], [] for text in texts: began = time.perf_counter() - # The path LLM2Vec.encode takes for one text, without its .to(device) that would undo the device map. features = model.tokenize([model.prepare_for_tokenization(model._convert_to_str("", text))]) - features = batch_to_device(features, input_device) with torch.no_grad(): vector = model.forward(features).detach().cpu().to(torch.float32)[0] assert vector.shape == (EMBEDDING_DIM,) vectors.append(vector.numpy().tobytes()) encode_seconds.append(time.perf_counter() - began) - Path(args.out).write_bytes(msgpack.packb({ + out.write_bytes(msgpack.packb({ "vectors": vectors, "encodeSeconds": encode_seconds, "loadSeconds": load_seconds, - "placement": placement, - "peakVramBytes": torch.cuda.max_memory_allocated(0) if torch.cuda.is_available() else 0, + "placement": {"cpuBytes": sum(p.numel() * p.element_size() for p in model.parameters())}, "peakRssBytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024, - "device": args.device, + "device": "cpu", + "assembly": ASSEMBLY, }, use_bin_type=True)) return 0 diff --git a/scripts/test_text_encoding.py b/scripts/test_text_encoding.py index 65aa1e3..899ffe6 100644 --- a/scripts/test_text_encoding.py +++ b/scripts/test_text_encoding.py @@ -27,7 +27,7 @@ class FakeChild: return { "vectors": [vector(len(text)) for text in texts], "encodeSeconds": [0.5] * len(texts), - "placement": {"gpuBytes": 11 * GIB, "cpuBytes": 5 * GIB}, + "placement": {"cpuBytes": 5 * GIB}, } @@ -36,28 +36,6 @@ def open_encoder(directory, child, identity="id-a"): return te.CachedTextEncoder(store, identity, child) -class PlanPlacement(unittest.TestCase): - def test_refuses_below_two_gib_with_a_typed_error_and_no_cpu_placement(self): - for free in (0, 1 * GIB, int(1.5 * GIB), 2 * GIB - 1): - with self.assertRaises(te.GpuBusy) as raised: - te.plan_placement(free, 64 * GIB) - self.assertEqual(raised.exception.code, "gpu-busy") - self.assertEqual(raised.exception.free_vram, free) - - def test_splits_free_vram_and_ram_each_less_the_reserve(self): - self.assertEqual(te.plan_placement(12 * GIB, 7 * GIB), {0: 12 * GIB - te.GPU_RESERVE, "cpu": 7 * GIB - te.CPU_RESERVE}) - self.assertEqual(te.plan_placement(14 * GIB, 10 * GIB), {0: 14 * GIB - te.GPU_RESERVE, "cpu": 10 * GIB - te.CPU_RESERVE}) - - def test_the_gpu_reserve_exceeds_the_old_one_that_crashed_at_the_card_ceiling(self): - self.assertGreater(te.GPU_RESERVE, int(1.5 * GIB)) - - def test_exactly_two_gib_free_is_admitted(self): - self.assertEqual(te.plan_placement(2 * GIB, 10 * GIB)[0], 2 * GIB - te.GPU_RESERVE) - - def test_ram_below_the_reserve_never_goes_negative(self): - self.assertEqual(te.plan_placement(12 * GIB, GIB)["cpu"], 0) - - class Identity(unittest.TestCase): def test_key_changes_with_identity_and_text_and_is_stable(self): self.assertEqual(te.embedding_key("a", "x"), te.embedding_key("a", "x")) @@ -65,12 +43,26 @@ class Identity(unittest.TestCase): self.assertNotEqual(te.embedding_key("a", "x"), te.embedding_key("a", "y")) self.assertNotEqual(te.embedding_key("ab", "x"), te.embedding_key("a", "bx")) - def test_identity_names_every_input(self): - base = dict(revisions={"r/one": "111", "r/two": "222"}, dtype="bfloat16", transformers_version="5.1.0", peft_version="0.18.1", kimodo_commit="6bb5848") - identity = te.encoder_identity(**base) - for field, other in (("dtype", "float32"), ("transformers_version", "5.2.0"), ("peft_version", "0.19.0"), ("kimodo_commit", "abcdef0"), - ("revisions", {"r/one": "111", "r/two": "333"})): - self.assertNotEqual(identity, te.encoder_identity(**{**base, field: other}), field) + def test_every_identity_field_changes_the_key(self): + # The repos' names alone exceed 40 characters, so a key built from a truncated identity cannot tell + # the later fields apart: each field here sits past that point. + base = dict( + assembly=te.ASSEMBLY, device="cpu", revisions={"r/one": "111", "r/two": "222", "r/three": "333"}, dtype="bfloat16", + torch_version="2.5.1+cpu", transformers_version="5.1.0", peft_version="0.18.1", accelerate_version="1.13.0", + kimodo_commit="6bb5848", loader_sha256="a" * 64, + ) + key = te.embedding_key(te.encoder_identity(**base), "x") + changes = dict( + assembly="llm2vec-A", device="other-device", dtype="float32", torch_version="2.6.0+cpu", transformers_version="5.2.0", + peft_version="0.19.0", accelerate_version="1.14.0", kimodo_commit="abcdef0", loader_sha256="b" * 64, + ) + for repo in base["revisions"]: + changes[f"revision:{repo}"] = {**base["revisions"], repo: "999"} + self.assertEqual(len(changes), 12) + for field, other in changes.items(): + name = "revisions" if field.startswith("revision:") else field + changed = te.embedding_key(te.encoder_identity(**{**base, name: other}), "x") + self.assertNotEqual(key, changed, field) def test_snapshot_revision_is_read_from_the_mounted_cache(self): with tempfile.TemporaryDirectory() as home: @@ -144,17 +136,6 @@ class CachedEncoder(unittest.TestCase): self.assertEqual(tuple(features.shape), (1, 1, 4096)) self.assertEqual(encoder.store.get_all(te.TEXT_EMBEDDING)[0].dim, 4096) - def test_gpu_busy_propagates_without_a_cpu_path_and_stores_nothing(self): - def busy(texts): - raise te.GpuBusy(GIB) - - with tempfile.TemporaryDirectory() as directory: - encoder = open_encoder(directory, busy) - with self.assertRaises(te.GpuBusy) as raised: - encoder([SECRET]) - self.assertNotIn(SECRET, str(raised.exception)) - self.assertEqual(encoder.count(), 0) - def test_the_contract_methods_return_self_or_the_device(self): with tempfile.TemporaryDirectory() as directory: encoder = open_encoder(directory, FakeChild()) @@ -163,6 +144,32 @@ class CachedEncoder(unittest.TestCase): self.assertEqual(encoder.get_device(), "cpu") +class ThreadCap(unittest.TestCase): + def test_accepts_one_to_sixteen(self): + for raw in ("1", "8", "16"): + self.assertEqual(te.thread_cap({"VILI_THREADS": raw}), int(raw)) + + def test_refuses_the_rest_without_echoing_the_value(self): + for raw in ("", "0", "17", "-1", "8.5", "eight", " 8", "\u0668"): + with self.assertRaises(ValueError) as raised: + te.thread_cap({"VILI_THREADS": raw}) + self.assertNotIn(raw or "NUL", str(raised.exception)) + with self.assertRaises(ValueError): + te.thread_cap({}) + + def test_apply_sets_torch_to_the_cap(self): + before = te.torch.get_num_threads() + try: + self.assertEqual(te.apply_thread_cap({"VILI_THREADS": "3"}), 3) + self.assertEqual(te.torch.get_num_threads(), 3) + finally: + te.torch.set_num_threads(before) + + def test_the_child_applies_the_cap_before_it_loads_anything(self): + source = (Path(te.__file__).with_name("kimodo-text-encode.py")).read_text() + self.assertLess(source.index("apply_thread_cap(os.environ)"), source.index("load_encoder(preset")) + + class ChildBoundary(unittest.TestCase): def run_stub(self, body, texts=("aa",)): with tempfile.TemporaryDirectory() as directory: @@ -170,16 +177,6 @@ class ChildBoundary(unittest.TestCase): stub.write_text(textwrap.dedent(body)) return te.run_child(list(texts), Path(directory) / "out.mc", child_path=stub, timeout=60) - def test_exit_three_is_gpu_busy_with_the_measured_free_vram(self): - with self.assertRaises(te.GpuBusy) as raised: - self.run_stub(""" - import sys, msgpack - out = sys.argv[sys.argv.index("--out") + 1] - open(out, "wb").write(msgpack.packb({"error": "gpu-busy", "freeVramBytes": 1234}, use_bin_type=True)) - sys.exit(3) - """) - self.assertEqual(raised.exception.free_vram, 1234) - def test_other_failures_name_the_exit_code_and_never_the_text(self): with self.assertRaises(RuntimeError) as raised: self.run_stub("import sys; sys.exit(1)", texts=(SECRET,)) diff --git a/scripts/text_encoding.py b/scripts/text_encoding.py index 668e529..26676cd 100644 --- a/scripts/text_encoding.py +++ b/scripts/text_encoding.py @@ -22,31 +22,26 @@ import torch from cultcache_py import CultCache, SingleFileMessagePackBackingStore, define_database_entry_type, database_entry_field, leaves as t GIB = 1024 ** 3 -MIN_FREE_VRAM = 2 * GIB -GPU_RESERVE = int(3 * GIB) # headroom for activations and offloaded layers streamed through the card -CPU_RESERVE = int(1.5 * GIB) EMBEDDING_DIM = 4096 EMBEDDING_BYTES = EMBEDDING_DIM * 4 -CHILD_EXIT_GPU_BUSY = 3 CHILD_PATH = Path(__file__).with_name("kimodo-text-encode.py") -LLAMA_BASE = "meta-llama/Meta-Llama-3-8B-Instruct" # the base model the mntp adapter repo points at +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 -class GpuBusy(RuntimeError): - """The card has too little free VRAM to host the encoder; there is no CPU fallback.""" - - code = "gpu-busy" - - def __init__(self, free_vram: int) -> None: - self.free_vram = free_vram - super().__init__(f"gpu-busy: {free_vram / GIB:.2f} GiB of VRAM free, {MIN_FREE_VRAM / GIB:.0f} GiB needed") +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 plan_placement(free_vram: int, mem_available: int) -> dict: - """The accelerate max_memory for one child call: free VRAM and available RAM, each less its reserve.""" - if free_vram < MIN_FREE_VRAM: - raise GpuBusy(free_vram) - return {0: free_vram - GPU_RESERVE, "cpu": max(mem_available - CPU_RESERVE, 0)} +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: @@ -55,10 +50,16 @@ def read_snapshot_revision(hf_home: Path, repo_id: str) -> str: return ref.read_text().strip() -def encoder_identity(*, revisions: dict, dtype: str, transformers_version: str, peft_version: str, kimodo_commit: str) -> str: - """The canonical string a vector is keyed under: every input that can change it.""" +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"dtype={dtype}", f"transformers={transformers_version}", f"peft={peft_version}", f"kimodo={kimodo_commit}"] + 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) @@ -110,7 +111,7 @@ def open_store(path) -> CultCache: 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 GpuBusy or RuntimeError.""" + """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"), @@ -118,8 +119,6 @@ def run_child(texts: list, out_path, child_path=CHILD_PATH, timeout: float = 180 timeout=timeout, ) answer = msgpack.unpackb(Path(out_path).read_bytes(), raw=False) if Path(out_path).exists() else {} - if process.returncode == CHILD_EXIT_GPU_BUSY: - raise GpuBusy(int(answer.get("freeVramBytes", 0))) if process.returncode != 0: raise RuntimeError(f"text encoder child failed with exit code {process.returncode}") return answer @@ -196,7 +195,7 @@ class CachedTextEncoder: dim=EMBEDDING_DIM, vector=vector, vectorSha256=hashlib.sha256(vector).hexdigest(), - gpuBytes=int(placement.get("gpuBytes", 0)), + gpuBytes=0, cpuBytes=int(placement.get("cpuBytes", 0)), encodeSeconds=float(took), createdAt=datetime.now(timezone.utc).isoformat(), diff --git a/vendor/CultLib b/vendor/CultLib new file mode 160000 index 0000000..02cf04b --- /dev/null +++ b/vendor/CultLib @@ -0,0 +1 @@ +Subproject commit 02cf04b5c9266772947a69a6c25218e0d157507d