diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 07c9b3e68a..fd49d9622b 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -1695,7 +1695,9 @@ encoder's card needs code in the ready-made encoder (RTX 20/30/40/50). `"vram_elastic": true`; [VRAM_ELASTIC.md](VRAM_ELASTIC.md#the-image-encoder-on-demand)):** the GPU encoder does not start with the server, so the expert cache keeps all of the VRAM. The first picture's turn makes the cache keep `"elastic_mib"` (default 2048) free and starts the encoder there (a few seconds); later pictures use it as it is, and -after `"elastic_idle_s"` (default 120) without a picture it stops and the cache takes the VRAM back. Measured on an R9700 +after `"elastic_idle_s"` (default 120) without a picture it stops and the cache takes the VRAM back. A turn that +re-sends pictures the hash cache already knows needs no encoder at all - its files survive the stop - so it takes +neither the turn nor a start; only a picture the cache lacks wakes it again. Measured on an R9700 32 GB (gfx1201, Linux, UD-IQ4_XS, the BF16 encoder at up to 1,024 tokens, #2189 for HIP's virtual memory): the encoder took 1,339 MiB and started in 4.7 s (2.7 s the next time); a request with one picture took 8.6 s with the start and 5.0 s on a later cold start, and 2 / 5 / 10 pictures with the encoder running took 3.8 / 5.6 / 9.0 s (~0.45 s to encode diff --git a/docs/VRAM_ELASTIC.md b/docs/VRAM_ELASTIC.md index 64ea795e28..7a1a11633d 100644 --- a/docs/VRAM_ELASTIC.md +++ b/docs/VRAM_ELASTIC.md @@ -114,6 +114,10 @@ and starts the encoder with the model, as before) the GPU encoder is not started `VRAM` command first keeps `"elastic_mib"` (default 2048) free on top of the reserve - the last `POST /v1/vram`, or 0 when none was made - and then the encoder starts. If the cache cannot shrink (the resident low-RAM mode, an engine without `--vram-elastic`) it starts anyway, as a non-elastic encoder would, and the log says why. +- **Pictures the hash cache already knows never wake it**: the cache and its `.sve` files survive a stop, and most + clients re-send the conversation's pictures on every turn - such a turn (every picture a hit, elastic only) takes + neither the FIFO turn nor a start **(code; `serve/test_vision_cache_wake.py`)**. A turn with one picture the cache + lacks takes both, for all of its pictures. - **Idle** for `"elastic_idle_s"` (default 120) without a picture: a thread stops the encoder between requests and sends `VRAM` with that reserve again, so the cache grows back to the size it had (0, not the start's reserve: that would leave it a segment or more smaller). diff --git a/serve/server.py b/serve/server.py index 5a119ec7cf..51dad3a841 100644 --- a/serve/server.py +++ b/serve/server.py @@ -2696,6 +2696,26 @@ def normalize(data: bytes) -> bytes: im.save(out, format="PNG") return out.getvalue() + def key_of(self, source) -> str: + """The encoder-cache key of an image: the hash encode() stores under (load+normalize first, so the same + picture in another encoding or size hits too).""" + data = self.normalize(source if isinstance(source, bytes) else self.load(source)) + return hashlib.sha256(data).hexdigest()[:32] + + def cached(self, source) -> tuple[Path, int] | None: + """(embeddings file, image tokens) if this image was already encoded, else None. The cache and the .sve + files survive an elastic unload (see unload()), so an all-cached conversation never needs the encoder up.""" + try: + key = self.key_of(source) + except Exception: + return None # unreadable source: let encode() produce the proper error later + with self.lock: + entry = self.cache.get(key) + if entry is not None: + self.cache[key] = self.cache.pop(key) # most recently used last, as encode() does on a hit: the + return entry # re-sent pictures this answers for are the ones that must + # not be the first evicted when the 64 entries fill up + def encode_all(self, sources: list, cancel=None) -> list[tuple[Path, int]]: """encode() for every image of one request: the files of this request's earlier images are never evicted to make room for a later one (#1072: with more than 64 unique images the first was deleted before the request @@ -4338,17 +4358,25 @@ def fetch(src): return fetched[src] if src in fetched else Vision.download(src) \ if src.startswith(("http://", "https://")) else src images = [fetch(src) for src in images] - # Encode only while the engine is idle: the engine and the image encoder (a separate process) must not - # run on the GPU at the same time - an encode during a running request left that request stuck at - # "reading the prompt" with CPU and GPU busy, for good (reproduced). So encoding takes its turn in the - # same FIFO as the requests. - with self._turn(cancel): - self._vision_up() - try: - encoded = (self.vision.encode_all(images, cancel) if hasattr(self.vision, "encode_all") - else [self.vision.encode(src) for src in images]) - finally: - self.last_image_at = time.time() + # Pictures the encoder already know never wake it: the hash cache and its .sve files survive an elastic + # unload, and most clients re-send every picture of the conversation on every turn - a text-only turn + # would otherwise pay the elastic start (a few seconds + the VRAM the expert cache gives up) on a hit. + cached = ([self.vision.cached(src) for src in images] if getattr(self.vision, "elastic", False) + and hasattr(self.vision, "cached") else [None] * len(images)) + if all(c is not None for c in cached): + encoded = cached + else: + # Encode only while the engine is idle: the engine and the image encoder (a separate process) must not + # run on the GPU at the same time - an encode during a running request left that request stuck at + # "reading the prompt" with CPU and GPU busy, for good (reproduced). So encoding takes its turn in the + # same FIFO as the requests. + with self._turn(cancel): + self._vision_up() + try: + encoded = (self.vision.encode_all(images, cancel) if hasattr(self.vision, "encode_all") + else [self.vision.encode(src) for src in images]) + finally: + self.last_image_at = time.time() # one <|image_pad|> per image -> one per image token. Only the markers the template writes for an image # (right after <|vision_start|>) are images: the same text inside a message (an agent reading these docs, # #150) is kept as plain text, or it took an image's place and the counts no longer matched. diff --git a/serve/test_vision_cache_wake.py b/serve/test_vision_cache_wake.py new file mode 100644 index 0000000000..8ee305bc10 --- /dev/null +++ b/serve/test_vision_cache_wake.py @@ -0,0 +1,138 @@ +"""The elastic encoder stays down for pictures it has already encoded: the hash cache and its .sve files survive an +unload (Vision.unload keeps them), so a conversation that re-sends its pictures on every turn - which is what most +clients do - must not pay the elastic start (seconds + the VRAM the expert cache gives up) on a cache hit. A new +picture, alone or alongside cached ones, still wakes it.""" +import io +import tempfile +import unittest +from pathlib import Path + +from PIL import Image + +from serve.frontend import ChatTemplate +from serve.server import ByteTokenizer, Service, Vision +from serve.test_vram import VramEngine + +ROOT = Path(__file__).resolve().parents[1] + + +class CachingEncoder: + """An encoder process stand-in with Vision's cache semantics: every picture encodes to three image tokens and + keeps its (file, count) in .cache, which - as on a real unload - survives restarts and unloads.""" + + def __init__(self, d): + self.elastic, self.elastic_mib, self.elastic_idle_s = True, 2048, 60.0 + self.up, self.starts, self.unloads, self.encodes = False, 0, 0, 0 + self.dir = Path(d) + self.cache = {} + + def alive(self): + return self.up + + def restart(self): + self.up, self.starts = True, self.starts + 1 + + def unload(self): + self.up, self.unloads = False, self.unloads + 1 + + def key_of(self, source): + return str(source) + + def cached(self, source): + return self.cache.get(self.key_of(source)) + + def encode(self, source, keep=()): + assert self.up, "a picture reached an encoder that was not started" + self.encodes += 1 + row = self.dir / f"{source}.sve" + row.write_bytes(b"rows") + self.cache[self.key_of(source)] = (row, 3) + return self.cache[self.key_of(source)] + + def encode_all(self, sources, cancel=None): + return [self.encode(s) for s in sources] + + +class CacheWake(unittest.TestCase): + def setUp(self): + self.tok = ByteTokenizer() + self.engine = VramEngine(self.tok) + self.engine.loaded = True + self.vision = CachingEncoder(tempfile.mkdtemp()) + self.svc = Service(self.engine, self.tok, ChatTemplate(ROOT / "serve/chat_template.jinja"), vision=self.vision) + self.addCleanup(self.drop_embeddings) + + def drop_embeddings(self): + path = getattr(getattr(self.svc, "embeddings", None), "path", None) + if path is not None: + path.unlink(missing_ok=True) + + def send(self, *sources): + content = [{"type": "text", "text": "what is it"}] + [{"type": "image", "source": s} for s in sources] + return self.svc.prepare([{"role": "user", "content": content}], None, {}) + + def test_a_re_sent_picture_does_not_wake_the_idle_encoder(self): + ids1 = self.send("a.png") + self.assertEqual(self.vision.starts, 1) # a new picture: the normal start + self.svc.last_image_at = 0 + self.assertEqual(self.svc._vision_idle_down(), "down") # the 60 s idle passed + ids2 = self.send("a.png") # the next turn re-sends the history + self.assertEqual(self.vision.starts, 1) # already encoded: no start + self.assertFalse(self.vision.up) + self.assertFalse(self.svc.vision_held) + self.assertEqual(self.engine.calls, [2048, 0]) # and no fresh share request either + self.assertEqual(ids2, ids1) # the prompt is built the same, from the cache + + def test_a_new_picture_alongside_cached_ones_still_wakes_it(self): + self.send("a.png") + self.svc.last_image_at = 0 + self.svc._vision_idle_down() + self.send("a.png", "b.png") # one hit, one miss + self.assertEqual(self.vision.starts, 2) + self.assertTrue(self.vision.up) + + def test_a_cache_hit_while_the_encoder_runs_changes_nothing(self): + self.send("a.png") + self.send("a.png") # the encoder is up: nothing to reserve, start or encode + self.assertEqual(self.vision.starts, 1) + self.assertEqual(self.engine.calls, [2048]) + + def test_the_cache_lookup_needs_no_encoder_process(self): + v = Vision({"exe": "e", "mmproj": "m", "model": "t", "elastic": True}, lazy=True) + self.assertFalse(v.alive()) # nothing started + img = Image.new("RGB", (8, 6), (9, 9, 9)) + buf = io.BytesIO() + img.save(buf, format="PNG") + png = buf.getvalue() + self.assertIsNone(v.cached(png)) # cold: nobody encoded it + v.cache[v.key_of(png)] = (Path("x.sve"), 42) # what encode() would have stored + self.assertEqual(v.cached(png), (Path("x.sve"), 42)) # warm: the file alone answers, process down + self.assertIsNone(v.cached(b"not an image")) # unreadable: None - encode() gives that error + + def test_a_cache_hit_keeps_the_picture_from_aging_out(self): + v = Vision({"exe": "e", "mmproj": "m", "model": "t", "elastic": True}, lazy=True) + pngs = [] + for shade in range(3): # three distinct pictures + img = Image.new("RGB", (8, 6), (shade, shade, shade)) + buf = io.BytesIO() + img.save(buf, format="PNG") + pngs.append(buf.getvalue()) + for png in pngs: # encoded oldest to newest, as encode() would cache + v.cache[v.key_of(png)] = (Path(f"{len(v.cache)}.sve"), 3) + self.assertEqual(v.cached(pngs[0]), (Path("0.sve"), 3)) # the oldest is asked for again (re-sent history) + self.assertEqual(list(v.cache)[-1], v.key_of(pngs[0])) # the lookup marks it most recently used, so + # the 64-entry eviction (encode, oldest first) + # takes a picture nobody re-sent instead + + def test_a_non_elastic_turn_takes_todays_path_even_all_cached(self): + self.vision.elastic = False # a plain always-on encoder + self.vision.up = True # it started with the server, not with a picture + self.send("a.png") # encode_all ran (a cache hit inside encode()) + self.assertEqual(self.vision.encodes, 1) + self.send("a.png") # the next turn: today's path unchanged, hits and all + self.assertEqual(self.vision.encodes, 2) + self.assertEqual(self.engine.calls, []) # no make-room request either: the elastic machinery stays out + + +if __name__ == "__main__": + unittest.main()