From 44c39f9951722c7b847c9d4278c120b1e2be9683 Mon Sep 17 00:00:00 2001 From: Owen Qwen Date: Sat, 12 Sep 2026 10:24:32 -0500 Subject: [PATCH] feat: add persistent conversation checkpoints --- README.md | 7 + docs/benchmarks/conversation-cache-m5.json | 265 ++++++++++++++ docs/conversation-cache-results.md | 101 +++++ docs/conversation-cache.md | 121 ++++++ scripts/bench_conversation_cache.py | 120 ++++++ src/edge0/backends/mlx/checkpoint.py | 87 +++++ src/edge0/cli.py | 23 ++ src/edge0/conversation/__init__.py | 4 + src/edge0/conversation/store.py | 407 +++++++++++++++++++++ src/edge0/engine/base.py | 96 ++++- src/edge0/engine/checkpoint.py | 50 +++ src/edge0/engine/ling.py | 13 + src/edge0/engine/qwen.py | 8 + src/edge0/models/base.py | 2 + src/edge0/server/app.py | 2 + src/edge0/server/chat.py | 25 +- tests/test_cache_continuation.py | 229 ++++++++++++ tests/test_conversation_cache.py | 210 +++++++++++ 18 files changed, 1764 insertions(+), 6 deletions(-) create mode 100644 docs/benchmarks/conversation-cache-m5.json create mode 100644 docs/conversation-cache-results.md create mode 100644 docs/conversation-cache.md create mode 100644 scripts/bench_conversation_cache.py create mode 100644 src/edge0/backends/mlx/checkpoint.py create mode 100644 src/edge0/conversation/__init__.py create mode 100644 src/edge0/conversation/store.py create mode 100644 src/edge0/engine/checkpoint.py create mode 100644 tests/test_cache_continuation.py create mode 100644 tests/test_conversation_cache.py diff --git a/README.md b/README.md index fb8168a..d4be59a 100644 --- a/README.md +++ b/README.md @@ -266,3 +266,10 @@ examples/demo.py # minimal API walkthrough ## License Apache-2.0, including vendored third-party code (see [NOTICE](NOTICE)). + +### Persistent conversation caching (opt-in) + +Reuse processed prompt prefixes across requests and process restarts with +`edge0 serve /path/to/model --cache-dir /path/to/conversation-cache`. +Defaults: 20 GiB of checkpoint payloads and checkpoints every 2,048 tokens. +See [configuration, guarantees, and benchmarks](docs/conversation-cache.md). diff --git a/docs/benchmarks/conversation-cache-m5.json b/docs/benchmarks/conversation-cache-m5.json new file mode 100644 index 0000000..e1aa73e --- /dev/null +++ b/docs/benchmarks/conversation-cache-m5.json @@ -0,0 +1,265 @@ +[ + { + "case": "disabled", + "output_tokens": [ + 1092, + 13321, + 300, + 268, + 25825, + 198, + 198, + 678 + ], + "ttft_s": 6.5136767080002755, + "wall_s": 7.46466375, + "post_first_token_effective_tok_s": 7.360773271190409, + "tokenization_s": 0.018990666999343375, + "prefill_s": 6.318684333999954, + "decode_s": 1.1266685829996277, + "decode_tok_s": 7.100579638691002, + "peak_mlx_bytes": 2731136504, + "peak_rss_bytes": 4904845312, + "io_read_blocks": 0, + "io_write_blocks": 0, + "system_disk_read_bytes": 573726720, + "system_disk_write_bytes": 179904512, + "usage": { + "prompt_tokens": 2799, + "completion_tokens": 8, + "total_tokens": 2807 + } + }, + { + "case": "first_write", + "output_tokens": [ + 1092, + 13321, + 300, + 268, + 25825, + 198, + 198, + 678 + ], + "ttft_s": 6.774661375000505, + "wall_s": 7.7653235410007255, + "post_first_token_effective_tok_s": 7.0659809572241645, + "tokenization_s": 0.006916833000104816, + "prefill_s": 5.296447042001091, + "decode_s": 0.5170845000002373, + "decode_tok_s": 15.471359129883663, + "peak_mlx_bytes": 3387898722, + "peak_rss_bytes": 5075402752, + "io_read_blocks": 0, + "io_write_blocks": 0, + "system_disk_read_bytes": 260599808, + "system_disk_write_bytes": 1423396864, + "usage": { + "prompt_tokens": 2799, + "completion_tokens": 8, + "total_tokens": 2807, + "prompt_tokens_details": { + "cached_tokens": 0 + } + }, + "cache": { + "lookup_s": 0.0003477079999356647, + "writes": 3, + "write_s": 1.9365312069994616, + "stored_bytes": 761533603, + "prefill_s": 5.296447042001091, + "remaining_prefill_tokens": 2799, + "tokenization_s": 0.006814792000113812 + } + }, + { + "case": "repeat", + "output_tokens": [ + 1092, + 13321, + 300, + 268, + 25825, + 198, + 198, + 678 + ], + "ttft_s": 0.5033716249999998, + "wall_s": 1.0222226670002783, + "post_first_token_effective_tok_s": 13.491348062082608, + "tokenization_s": 0.0013498750004146132, + "prefill_s": 1.583999619469978e-06, + "decode_s": 0.6283079169998018, + "decode_tok_s": 12.732610529882155, + "peak_mlx_bytes": 1777306570, + "peak_rss_bytes": 4981604352, + "io_read_blocks": 0, + "io_write_blocks": 0, + "system_disk_read_bytes": 27803648, + "system_disk_write_bytes": 25305088, + "usage": { + "prompt_tokens": 2799, + "completion_tokens": 8, + "total_tokens": 2807, + "prompt_tokens_details": { + "cached_tokens": 2799 + } + }, + "cache": { + "lookup_s": 0.00028974999986530747, + "restore_s": 0.38771641599942086, + "reused_tokens": 2799, + "prefill_s": 1.583999619469978e-06, + "remaining_prefill_tokens": 0, + "tokenization_s": 0.0012944999998580897 + } + }, + { + "case": "appended_turn", + "output_tokens": [ + 198, + 198, + 6952, + 698, + 268, + 9379, + 3189, + 391 + ], + "ttft_s": 1.6190897500000574, + "wall_s": 2.5677901669996572, + "post_first_token_effective_tok_s": 7.378514728747033, + "tokenization_s": 0.005887457999961043, + "prefill_s": 0.5801039170000877, + "decode_s": 0.4516456660003314, + "decode_tok_s": 17.71300070439319, + "peak_mlx_bytes": 1819663332, + "peak_rss_bytes": 4993744896, + "io_read_blocks": 0, + "io_write_blocks": 0, + "system_disk_read_bytes": 66650112, + "system_disk_write_bytes": 782823424, + "usage": { + "prompt_tokens": 2832, + "completion_tokens": 8, + "total_tokens": 2840, + "prompt_tokens_details": { + "cached_tokens": 2807 + } + }, + "cache": { + "lookup_s": 0.0002899589999287855, + "restore_s": 0.3203155000001061, + "reused_tokens": 2807, + "writes": 2, + "write_s": 1.206859083999916, + "stored_bytes": 992792613, + "prefill_s": 0.5801039170000877, + "remaining_prefill_tokens": 25, + "tokenization_s": 0.005829499999890686 + } + }, + { + "case": "branch", + "output_tokens": [ + 198, + 198, + 6952, + 698, + 1099, + 297, + 6474, + 3832 + ], + "ttft_s": 1.5322595829993588, + "wall_s": 2.667079666999598, + "post_first_token_effective_tok_s": 6.168378669616958, + "tokenization_s": 0.005418582999482169, + "prefill_s": 0.46757583300041006, + "decode_s": 0.6374825420007255, + "decode_tok_s": 12.549363273372427, + "peak_mlx_bytes": 1840880720, + "peak_rss_bytes": 5205737472, + "io_read_blocks": 0, + "io_write_blocks": 0, + "system_disk_read_bytes": 78704640, + "system_disk_write_bytes": 906326016, + "usage": { + "prompt_tokens": 2831, + "completion_tokens": 8, + "total_tokens": 2839, + "prompt_tokens_details": { + "cached_tokens": 2807 + } + }, + "cache": { + "lookup_s": 0.0002689579996513203, + "restore_s": 0.3277814999992188, + "reused_tokens": 2807, + "writes": 2, + "write_s": 1.2258794170002147, + "stored_bytes": 1223715239, + "prefill_s": 0.46757583300041006, + "remaining_prefill_tokens": 24, + "tokenization_s": 0.005361083000025246 + } + }, + { + "case": "process_restart", + "output_tokens": [ + 1092, + 13321, + 300, + 268, + 25825, + 198, + 198, + 678 + ], + "ttft_s": 0.6099199159998534, + "wall_s": 1.224041457999192, + "post_first_token_effective_tok_s": 11.398395140497348, + "tokenization_s": 0.005186958000194863, + "prefill_s": 6.250011210795492e-07, + "decode_s": 0.8147992079993855, + "decode_tok_s": 9.818369877460697, + "peak_mlx_bytes": 1190999148, + "peak_rss_bytes": 2118565888, + "io_read_blocks": 0, + "io_write_blocks": 0, + "system_disk_read_bytes": 38510592, + "system_disk_write_bytes": 36405248, + "usage": { + "prompt_tokens": 2799, + "completion_tokens": 8, + "total_tokens": 2807, + "prompt_tokens_details": { + "cached_tokens": 2799 + } + }, + "cache": { + "lookup_s": 0.003613083999880473, + "restore_s": 0.3949573749996489, + "reused_tokens": 2799, + "prefill_s": 6.250011210795492e-07, + "remaining_prefill_tokens": 0, + "tokenization_s": 0.005138458000146784 + } + }, + { + "case": "radix_lookup", + "checkpoints": 1000, + "lookup_us": 1.5507749999414955 + }, + { + "case": "radix_lookup", + "checkpoints": 10000, + "lookup_us": 1.5693625000494649 + }, + { + "case": "radix_lookup", + "checkpoints": 100000, + "lookup_us": 1.5628167000613757 + } +] diff --git a/docs/conversation-cache-results.md b/docs/conversation-cache-results.md new file mode 100644 index 0000000..561b552 --- /dev/null +++ b/docs/conversation-cache-results.md @@ -0,0 +1,101 @@ +# Local cache measurements + +Measured September 12, 2026 on an Apple M5 MacBook Air with 24 GB unified memory, +local `edge0-8b`, its shipped LoRA and prerouter, MLX 0.30.6, and mlx-lm 0.31.0. +The fixed coding fixture contains 2,799 prompt tokens. Each request generates eight +greedy tokens. The appended turn and branch contain 2,832 and 2,831 prompt tokens. +The interval is 2,048; the checkpoint budget is 20 GiB. + +These are sequential single samples on an active laptop, not confidence intervals. +The OS page cache was warm. The baseline runs first and therefore includes more +kernel/expert warmup. Initial model loading and artifact hashing are excluded from +request latency. A fresh process is measured separately; no claim of cold-SSD +performance is made. [Raw measurements](benchmarks/conversation-cache-m5.json) +include output token IDs, memory, decode, and disk counters. + +| Request | Reused tokens | TTFT (s) | Wall (s) | Restore (s) | Remaining prefill (s) | Writes (s) | +|---|---:|---:|---:|---:|---:|---:| +| Disabled | 0 | 6.514 | 7.465 | — | 6.319 | — | +| First write | 0 | 6.775 | 7.765 | — | 5.296 | 1.937 | +| Repeat | 2,799 | 0.503 | 1.022 | 0.388 | <0.001 | 0 | +| Appended turn | 2,807 | 1.619 | 2.568 | 0.320 | 0.580 | 1.207 | +| Branch | 2,807 | 1.532 | 2.667 | 0.328 | 0.468 | 1.226 | +| Process restart | 2,799 | 0.610 | 1.224 | 0.395 | <0.001 | 0 | + +Write time includes prompt and completed-generation snapshots, so not all of it +falls before the first token. Prompt and completion counts remain unchanged by +reuse. Disabled, first-write, repeat, and restarted requests produced identical +eight-token outputs. Appended/branched requests processed only 25/24 unmatched +tokens through prefill. + +At this measured prefix length, the first cached request cost 0.261 s extra TTFT +and 0.301 s extra total time. One repeat saved 6.010 s TTFT and 6.442 s total time +against the disabled sample, so the first repeat amortized the observed initial +cost. Even charging the full 1.937 s publication time as overhead, one repeat +covered it. This is a measured **reuse-count** break-even at 2,799 tokens, not a +measured minimum token-length threshold. Short-prompt break-even, randomized +request-order trials, sustained workloads, and cold filesystem trials remain +unmeasured. Earlier development samples ranged from 5.45–6.36 s disabled TTFT +and 0.49 s repeat TTFT; laptop load and warmup affect the absolute values. + +| Request | Peak MLX (GiB) | Sampled peak RSS (GiB) | Decode (tokens/s) | +|---|---:|---:|---:| +| Disabled | 2.54 | 4.57 | 7.10 | +| First write | 3.16 | 4.73 | 15.47 | +| Repeat | 1.66 | 4.64 | 12.73 | +| Appended turn | 1.69 | 4.65 | 17.71 | +| Branch | 1.71 | 4.85 | 12.55 | +| Process restart | 1.11 | 1.97 | 9.82 | + +Publication increased peak MLX allocation by about 24% in this sample. The active +attention cache still lives in RAM; caching is not active-context offloading. +RSS includes expert caches and other allocations, and is sampled every 20 ms. +Decode values cover only eight tokens and have substantial warmup/noise; they do +not establish a decode-speed improvement. Including final synchronous publication, +the effective rate after the first token was 7.07 tokens/s for first write versus +13.49 for repeat. + +Unique payloads occupied 761,533,603 bytes after the first conversation, then +1,223,715,239 bytes after both branches. System-wide disk counters reported about +1.42 GB written during first write, 0.78 GB for the appended turn, and 0.91 GB for +the branch. These include other processes and filesystem behavior. Per-process +block counters remained zero on this macOS run. Content deduplication saves +retained storage, but the current synchronous codec reserializes shared blocks; +it does not eliminate their write/checksum work. + +A separate radix-only benchmark performed 10,000 lookups on 130-token keys: + +| Checkpoints | Mean lookup (µs) | +|---:|---:| +| 1,000 | 1.55 | +| 10,000 | 1.57 | +| 100,000 | 1.56 | + +This isolates index traversal. It excludes SQLite startup/rebuild, lock contention, +and payload restore. End-to-end warm lookup was about 0.27–0.35 ms in the verified +run; first lookup after restart, including index construction, was 3.61 ms. + +## Correctness coverage and remaining limits + +The tests cover radix branches and exact hits, incremental local index updates, +namespace/artifact invalidation, reference cleanup and LRU eviction, physical +storage limits, missing/corrupt payloads and metadata, repair of shared corrupt +blocks, interrupted publication including abrupt process exit, and concurrent +threads/processes. Continuation coverage includes attention plus recurrent state, +family prerouter fields and expert staging, exact hits, one-token suffixes, +intermediate prefill boundaries, full sampling history, EOS, generation limits, +and callback cancellation. HTTP session tests verify released idle context and +unchanged total token accounting. + +Real 8B continuation logits matched within `rtol=1e-4, atol=1e-4`; restored prompt +logits used `1e-5`. The actual small Qwen gated-delta/attention backbone also passed +with random weights at `1e-5`. Real-weight **35B validation remains outstanding** +because that checkpoint is unavailable locally. + +A checkpoint resumes the saved execution boundary. Hybrid prerouter behavior can +depend on prefill versus decode execution and chunk boundaries; the continuation +contract is equality with the same uninterrupted saved trajectory, not arbitrary +re-chunking of a previously decoded conversation. Corrupt entries become misses; +whole-database destruction or cache filesystem failure is not a recovery mechanism +for the inference service itself. No background write queue or quantized KV format +is included in this phase. diff --git a/docs/conversation-cache.md b/docs/conversation-cache.md new file mode 100644 index 0000000..77b558a --- /dev/null +++ b/docs/conversation-cache.md @@ -0,0 +1,121 @@ +# Persistent conversation checkpoints + +Caching is opt-in. It saves repeated prompt computation and releases completed +HTTP requests' attention/recurrent state from RAM. It does not offload the active +attention working set, quantize KV, or implement Pi's tool protocol. + +```sh +edge0 serve /path/to/model --cache-dir /path/to/conversation-cache +edge0 chat /path/to/model --cache-dir /path/to/conversation-cache --prompt 'Hello' +edge0 cache inspect --cache-dir /path/to/conversation-cache +edge0 cache clear --cache-dir /path/to/conversation-cache +``` + +Use a separate directory from the model. Defaults are a 20 GiB checkpoint payload +budget and a 2,048-token interval. `--cache-budget-gib` and `--cache-interval` +override these. The separate exact-tokenization budget defaults to 16 MiB; +Python can configure it. SQLite's allocated pages are included in an additional +directory bound: checkpoint budget + tokenization budget + 64 KiB schema reserve. +Compaction and additional LRU eviction enforce that bound after writes. Atomic +publication and SQLite journaling require temporary disk headroom during a write. + +```python +from edge0 import AutoEngine +from edge0.conversation import CacheConfig + +engine = AutoEngine.from_pretrained( + '/path/to/model', name='edge0-8b', + conversation_cache=CacheConfig( + directory='/path/to/conversation-cache', + budget_bytes=20 * 1024**3, + interval=2048, + token_budget_bytes=16 * 1024**2, + ), +) +engine.reset() +tokens = engine.generate(prompt_ids, max_new_tokens=128) +print(engine.conversation_cache.metrics) +engine.reset() # release the active Python continuation when done +engine.close() +``` + +`generate` looks up a prefix when the engine position is zero. Direct Python +`prefill`/`step` callers retain their active continuation until `reset`. Server +requests still use the existing single generation lock; `ChatSession.run` resets +before lookup and releases request state on completion or cancellation. Expert +weight caches remain shared across requests. + +The store saves interval boundaries, prompt completion, and completed generation. +Only forwarded tokens are recorded: a sampled EOS that was never forwarded is +excluded. A cancellation leaves earlier published boundaries usable. An interval +checkpoint taken during prefill carries its lifecycle phase, so an exact hit +finishes prefill staging before decode. A one-token unmatched suffix uses the +prefill lifecycle. Sampling penalties receive the full prompt history. + +## Identity and persistence + +Namespace identity includes model and adapter contents, tokenizer and chat-template +artifacts, the effective tokenizer vocabulary/serialization, model configuration, +checkpoint interval, family/backend environment settings, backend versions, and +Edge0 source contents. Artifact SHA-256 hashes are computed at initialization and +reused on later starts only when device, inode, size, nanosecond mtime, and ctime +match. First initialization can therefore read the whole model. Do not mutate a +loaded engine's model, adapters, tokenizer, or inference configuration in place; +construct a new engine for those changes. + +A compact radix index searches token edges in memory, verifies full token equality, +and loads only the chosen boundary. Local writes update the index incrementally. +After another process changes the directory, the index is rebuilt from SQLite +metadata once; this synchronization cost is separate from steady-state radix +lookup. It does not load tensors or retain inactive conversation tensors. + +Attention keys and values use immutable content-addressed safetensors blocks. +Recurrent tensors, next-token logits, position, and family prerouter fields are +saved in a boundary-specific snapshot. Expert staging is rebuilt through family +hooks. No tensor dtype conversion, pickle, or executable-object deserialization +is used. Blocks deduplicate only when their serialized contents match exactly. +Currently shared blocks are reserialized and checksummed on publication; this +cost is included in the measured write overhead. + +Payload and metadata checksums turn corrupt entries into misses. A directory lock +coordinates readers, publication, recovery, and eviction across processes. Files +are fsynced and atomically renamed before transactional SQLite publication. +Unpublished managed files are collected on recovery. LRU checkpoint eviction drops +unreferenced blocks; oversized snapshots are skipped. Cleanup only removes managed +payload names. `clear` covers every model namespace in the selected cache directory. + +The exact-tokenization cache keys the entire message list, template namespace, +and thinking settings. It never joins independently tokenized messages. + +## Metrics and validation + +HTTP usage retains the full `prompt_tokens` count and adds +`prompt_tokens_details.cached_tokens`. The optional `edge0_cache` response field +(and terminal SSE chunk) reports lookup, tokenization, restore, remaining prefill, +publication, reused-token, eviction, and stored-byte metrics. The Python store +exposes the same dictionary. `write_s` includes synchronous serialization and +checksumming; `restore_s` includes checksum validation, tensor restore, and family +staging. Remaining `prefill_s` excludes recorded checkpoint writes. + +The benchmark uses a fixed 48-function Python review conversation, then an +appended turn and a branch. Run it from the repository with the installed environment: + +```sh +python scripts/bench_conversation_cache.py --model-dir /path/to/model \ + --cache-dir /path/to/benchmark-cache +python scripts/bench_conversation_cache.py --model-dir /path/to/model \ + --cache-dir /path/to/benchmark-cache --restart +python scripts/bench_conversation_cache.py --lookup-only +``` + +The first command clears the specified benchmark cache. Model initialization and +artifact hashing occur before request timers. The restart command launches a new +model process, but does **not** flush the OS filesystem cache. TTFT is measured at +the engine's first token callback, which follows its first decode forward. +`decode_tok_s` measures sampling and decode, excluding checkpoint writes; +`post_first_token_effective_tok_s` also includes final publication and request +cleanup. RSS is sampled every 20 ms; MLX reports its own peak allocation. +System disk counters include other processes, while per-process block counters +may remain zero on macOS. These counters are not a claim of isolated SSD traffic. + +See [measured results](conversation-cache-results.md) for local timings and limits. diff --git a/scripts/bench_conversation_cache.py b/scripts/bench_conversation_cache.py new file mode 100644 index 0000000..ab9cd48 --- /dev/null +++ b/scripts/bench_conversation_cache.py @@ -0,0 +1,120 @@ +"""Local reproducible coding-conversation benchmark; JSON output. + +Run once normally, then with --restart against the same cache directory. +Restart means a new Python/model process, not a cold OS filesystem cache. +""" +import argparse +import json +import resource +import threading +import time +from pathlib import Path + +import psutil +from edge0 import AutoEngine +from edge0.backends import core +from edge0.conversation import CacheConfig +from edge0.conversation.store import Radix +from edge0.server.chat import ChatMessage, ChatRequest, ChatSession + + +def fixture(): + code = '\n\n'.join( + f'def normalize_path_{i}(path: str) -> str:\n' + f' """Normalize a path for source module {i}."""\n' + ' parts = [part for part in path.split("/") if part and part != "."]\n' + ' return "/".join(parts)' + for i in range(48)) + return [ChatMessage('user', 'Review this Python module. Explain two concrete improvements.\n```python\n' + code + '\n```')] + + +def measure(engine, name, messages): + process = psutil.Process() + rss = [process.memory_info().rss] + stopped = threading.Event() + def watch(): + while not stopped.wait(0.02): + rss.append(process.memory_info().rss) + thread = threading.Thread(target=watch) + thread.start() + core.reset_peak_memory() + before = resource.getrusage(resource.RUSAGE_SELF) + disk_before = psutil.disk_io_counters() + start = time.perf_counter() + first = [] + def token(_): + if not first: + first.append(time.perf_counter() - start) + try: + tokens, meta = ChatSession(engine, ChatRequest(engine.name, messages, + temperature=0, max_tokens=8)).run(on_token=token) + finally: + stopped.set() + thread.join() + wall = time.perf_counter() - start + after = resource.getrusage(resource.RUSAGE_SELF) + disk_after = psutil.disk_io_counters() + result = dict(case=name, output_tokens=tokens, ttft_s=first[0] if first else None, wall_s=wall, + post_first_token_effective_tok_s=(len(tokens) - 1) / (wall - first[0]) if first and len(tokens) > 1 else None, + tokenization_s=engine.last_tokenization_s, + **engine.last_generation_metrics, + peak_mlx_bytes=core.get_peak_memory(), peak_rss_bytes=max(rss), + io_read_blocks=after.ru_inblock - before.ru_inblock, + io_write_blocks=after.ru_oublock - before.ru_oublock, + system_disk_read_bytes=disk_after.read_bytes - disk_before.read_bytes, + system_disk_write_bytes=disk_after.write_bytes - disk_before.write_bytes, + **{k: v for k, v in meta.items() if k != "wall_s"}) + print(json.dumps(result), flush=True) + return tokens + + +def lookup_benchmark(): + for count in (1000, 10000, 100000): + tree = Radix() + prefix = list(range(128)) + for i in range(count): + tree.insert(prefix + [i, i + 1], i) + started = time.perf_counter() + for i in range(10000): + assert tree.longest(prefix + [i % count, i % count + 1, -1]) == i % count + print(json.dumps(dict(case='radix_lookup', checkpoints=count, + lookup_us=(time.perf_counter() - started) * 100))) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--model-dir') + parser.add_argument('--cache-dir') + parser.add_argument('--restart', action='store_true') + parser.add_argument('--lookup-only', action='store_true') + args = parser.parse_args() + if args.lookup_only: + lookup_benchmark() + return + engine = AutoEngine.from_pretrained(args.model_dir, name='edge0-8b', + conversation_cache=CacheConfig(args.cache_dir)) + try: + messages = fixture() + if args.restart: + measure(engine, 'process_restart', messages) + return + cache = engine.conversation_cache + cache.clear() + engine.conversation_cache = None + baseline = measure(engine, 'disabled', messages) + engine.conversation_cache = cache + tokens = measure(engine, 'first_write', messages) + repeated = measure(engine, 'repeat', messages) + assert baseline == tokens == repeated, 'repeated prompt continuation changed' + appended = messages + [ChatMessage('assistant', engine._tok.decode(tokens)), + ChatMessage('user', 'Add type annotations and unit tests for the first function.')] + measure(engine, 'appended_turn', appended) + branch = messages + [ChatMessage('assistant', engine._tok.decode(tokens)), + ChatMessage('user', 'Instead, explain how to handle parent directory components.')] + measure(engine, 'branch', branch) + finally: + engine.close() + + +if __name__ == '__main__': + main() diff --git a/src/edge0/backends/mlx/checkpoint.py b/src/edge0/backends/mlx/checkpoint.py new file mode 100644 index 0000000..6247c45 --- /dev/null +++ b/src/edge0/backends/mlx/checkpoint.py @@ -0,0 +1,87 @@ +"""Lossless MLX safetensors codec. No dynamic classes or executable payloads.""" +import mlx.core as mx +from mlx_lm.models.cache import KVCache, ArraysCache + + +def write(engine, put, phase='ready'): + family_state = engine._checkpoint_family_state() + def nbytes(value): + if isinstance(value, mx.array): + return value.nbytes + if isinstance(value, dict): + return sum(nbytes(v) for v in value.values()) + if isinstance(value, (tuple, list)): + return sum(nbytes(v) for v in value) + return 0 + estimate = sum(c.nbytes for c in engine.cache) + nbytes(family_state) + nbytes(engine.next_logits()) + if estimate > engine.conversation_cache.config.budget_bytes: + from edge0.conversation.store import _Oversized + raise _Oversized('checkpoint exceeds disk budget') + tensors = {} + def encode(value): + if isinstance(value, mx.array): + key = str(len(tensors)) + tensors[key] = value + return {'tensor': key} + if isinstance(value, (list, tuple)): + return {'list': [encode(v) for v in value]} + if isinstance(value, dict): + return {'dict': [[encode(k), encode(v)] for k, v in value.items()]} + if value is None or isinstance(value, (str, int, float, bool)): + return value + raise TypeError(f'unsupported checkpoint value: {type(value)}') + layers = [] + for cache in engine.cache: + if type(cache) is KVCache: + blocks = [] + keys, values = cache.state + for start in range(0, cache.offset, engine.conversation_cache.config.interval): + end = min(start + engine.conversation_cache.config.interval, cache.offset) + data = {'keys': keys[..., start:end, :], 'values': values[..., start:end, :]} + blocks.append(put(lambda path, data=data: mx.save_safetensors(str(path), data))) + layers.append({'kind': 'kv', 'offset': cache.offset, 'blocks': blocks}) + elif type(cache) is ArraysCache: + layers.append({'kind': 'arrays', 'state': encode(cache.state), + 'left_padding': encode(cache.left_padding), 'lengths': encode(cache.lengths)}) + else: + raise ValueError(f'unsupported cache class {type(cache)}') + state = encode(family_state) + logits = encode(engine.next_logits()) + snapshot = put(lambda path: mx.save_safetensors(str(path), tensors)) + return dict(layers=layers, state=state, logits=logits, snapshot=snapshot, + pos=engine.pos, phase=phase) + + +def read(manifest, paths): + tensors = mx.load(str(paths[manifest['snapshot']])) + def decode(value): + if isinstance(value, dict): + if 'tensor' in value: + return tensors[value['tensor']] + if 'list' in value: + return [decode(v) for v in value['list']] + if 'dict' in value: + return {decode(k): decode(v) for k, v in value['dict']} + raise ValueError('invalid state') + return value + caches = [] + for layer in manifest['layers']: + if layer['kind'] == 'kv': + blocks = [mx.load(str(paths[k])) for k in layer['blocks']] + cache = KVCache() + cache.state = tuple(mx.concatenate([b[k] for b in blocks], axis=2) for k in ('keys', 'values')) + if cache.offset != layer['offset'] or cache.offset != manifest['pos']: + raise ValueError('invalid KV offset') + elif layer['kind'] == 'arrays': + state = decode(layer['state']) + cache = ArraysCache(len(state)) + cache.state = state + cache.left_padding = decode(layer['left_padding']) + cache.lengths = decode(layer['lengths']) + else: + raise ValueError('unknown layer cache') + caches.append(cache) + logits = decode(manifest['logits']) + state = decode(manifest['state']) + mx.eval(logits, *[c.state for c in caches]) + return caches, logits, state, manifest['pos'], manifest['phase'] diff --git a/src/edge0/cli.py b/src/edge0/cli.py index cefa632..5afea5c 100644 --- a/src/edge0/cli.py +++ b/src/edge0/cli.py @@ -66,6 +66,10 @@ def _engine_kwargs(args) -> dict: kw["prerouter"] = None if getattr(args, "no_lora", False): kw["lora"] = "" + if getattr(args, 'cache_dir', None): + from edge0.conversation import CacheConfig + kw['conversation_cache'] = CacheConfig(args.cache_dir, + int(args.cache_budget_gib * 1024**3), args.cache_interval) return kw @@ -211,6 +215,16 @@ def cmd_convert(args) -> int: return 0 +def cmd_cache(args): + import json + from edge0.conversation import CacheConfig, CheckpointStore + store = CheckpointStore(CacheConfig(args.cache_dir), '', maintenance=False) + if args.action == 'clear': + store.clear() + print(json.dumps(store.inspect(), indent=2)) + return 0 + + def main(argv: list[str] | None = None) -> int: ap = argparse.ArgumentParser(prog="edge0", description=__doc__) sub = ap.add_subparsers(dest="cmd", required=True) @@ -269,6 +283,15 @@ def main(argv: list[str] | None = None) -> int: help="one-shot legacy npz -> safetensors migration") p.set_defaults(fn=cmd_convert) + for command in ('demo', 'chat', 'serve'): + parser = sub.choices[command] + parser.add_argument('--cache-dir', default=None, help='enable persistent conversation caching') + parser.add_argument('--cache-budget-gib', type=float, default=20) + parser.add_argument('--cache-interval', type=int, default=2048) + parser = sub.add_parser('cache', help='inspect or clear a conversation cache') + parser.add_argument('action', choices=('inspect', 'clear')) + parser.add_argument('--cache-dir', required=True) + parser.set_defaults(fn=cmd_cache) args = ap.parse_args(argv) return args.fn(args) diff --git a/src/edge0/conversation/__init__.py b/src/edge0/conversation/__init__.py new file mode 100644 index 0000000..ae8b966 --- /dev/null +++ b/src/edge0/conversation/__init__.py @@ -0,0 +1,4 @@ +"""Opt-in persistent inference checkpoints (no tensor imports).""" +from .store import CacheConfig, CheckpointStore + +__all__ = ['CacheConfig', 'CheckpointStore'] diff --git a/src/edge0/conversation/store.py b/src/edge0/conversation/store.py new file mode 100644 index 0000000..683beaf --- /dev/null +++ b/src/edge0/conversation/store.py @@ -0,0 +1,407 @@ +"""Locked, transactional, content-addressed checkpoint storage. + +The index holds token edges only. Payloads are opened only for the selected +boundary; no inactive tensors are retained. SQLite is the publication point. +""" +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +import fcntl +import hashlib +import json +import os +import re +from pathlib import Path +import sqlite3 +import threading +import time +import uuid + + +def digest(data): + return hashlib.sha256(data).hexdigest() + + +def _json_default(value): + from enum import Enum + if isinstance(value, Enum): + return value.value + if isinstance(value, Path): + return str(value) + raise TypeError(f"unsupported namespace value: {type(value)}") + + +def packed(value): + return json.dumps(value, sort_keys=True, separators=(',', ':'), default=_json_default).encode() + + +@dataclass(frozen=True) +class CacheConfig: + directory: str = '' + budget_bytes: int = 20 * 1024**3 + interval: int = 2048 + token_budget_bytes: int = 16 * 1024**2 + + def __post_init__(self): + if self.budget_bytes < 0 or self.token_budget_bytes < 0 or self.interval < 1: + raise ValueError('cache budgets must be nonnegative and interval positive') + + +class Radix: + def __init__(self): + self.children = {} + self.value = None + + def insert(self, tokens, value): + node = self + tokens = tuple(tokens) + while tokens: + item = node.children.get(tokens[0]) + if item is None: + child = Radix() + child.value = value + node.children[tokens[0]] = (tokens, child) + return + edge, child = item + n = 0 + while n < min(len(edge), len(tokens)) and edge[n] == tokens[n]: + n += 1 + if n < len(edge): + split = Radix() + split.children[edge[n]] = (edge[n:], child) + node.children[tokens[0]] = (edge[:n], split) + child = split + node, tokens = child, tokens[n:] + node.value = value + + def remove(self, tokens): + node, pos, parents = self, 0, [] + while pos < len(tokens): + item = node.children.get(tokens[pos]) + if item is None: + return + edge, child = item + if tuple(tokens[pos:pos + len(edge)]) != edge: + return + parents.append((node, tokens[pos])) + pos += len(edge) + node = child + node.value = None + for parent, key in reversed(parents): + edge, child = parent.children[key] + if child.value is None and not child.children: + del parent.children[key] + elif child.value is None and len(child.children) == 1: + tail, grandchild = next(iter(child.children.values())) + parent.children[key] = (edge + tail, grandchild) + + def longest(self, tokens): + node, pos, best = self, 0, None + while pos < len(tokens): + item = node.children.get(tokens[pos]) + if item is None: + break + edge, child = item + if tuple(tokens[pos:pos + len(edge)]) != edge: + break + pos += len(edge) + node = child + if node.value is not None: + best = node.value + return best + + +class _Oversized(ValueError): + pass + + +def _managed_payload(path): + return bool(re.fullmatch(r'[0-9a-f]{64}\.safetensors|[0-9a-f]{32}\.tmp(?:\.safetensors)?', path.name)) + + +class CheckpointStore: + def __init__(self, config: CacheConfig, namespace: str, *, maintenance=True): + self.config, self.namespace = config, namespace + self.maintenance = maintenance + self.root = Path(config.directory).expanduser().resolve() + self.root.mkdir(parents=True, exist_ok=True) + self._mutex = threading.RLock() + self._version = None + self.index = Radix() + self.metrics = {} + with self.locked(enforce=maintenance) as db: + db.executescript(''' + CREATE TABLE IF NOT EXISTS checkpoints ( + id TEXT PRIMARY KEY, ns TEXT, tokens BLOB, manifest BLOB, used REAL); + CREATE TABLE IF NOT EXISTS objects (id TEXT PRIMARY KEY, size INTEGER); + CREATE TABLE IF NOT EXISTS refs (checkpoint TEXT, object TEXT, + PRIMARY KEY(checkpoint, object)); + CREATE TABLE IF NOT EXISTS tokenizations ( + id TEXT PRIMARY KEY, value BLOB, used REAL); + CREATE TABLE IF NOT EXISTS fingerprints ( + path TEXT PRIMARY KEY, identity TEXT, hash TEXT); + CREATE TABLE IF NOT EXISTS revision (value INTEGER); + INSERT INTO revision SELECT 0 WHERE NOT EXISTS (SELECT 1 FROM revision); + ''') + for table in ('checkpoints', 'tokenizations'): + if 'checksum' not in {row[1] for row in db.execute(f'PRAGMA table_info({table})')}: + db.execute(f'ALTER TABLE {table} ADD COLUMN checksum TEXT') + if maintenance: + self._cleanup(db) + self._evict(db) + + @contextmanager + def locked(self, *, enforce=False): + with self._mutex, open(self.root / 'cache.lock', 'a+b') as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + db = sqlite3.connect(self.root / 'metadata.sqlite') + try: + with db: + yield db + if self.maintenance and enforce: + self._bound_directory(db) + finally: + db.close() + fcntl.flock(lock, fcntl.LOCK_UN) + + def _bound_directory(self, db): + # Include allocated SQLite pages, not just live tensor bytes. The + # separate tokenization allowance and a fixed 64 KiB schema reserve + # keep even tiny test budgets usable. Compact only under pressure. + limit = self.config.budget_bytes + self.config.token_budget_bytes + 65536 + def physical(): + return sum(p.stat().st_size for p in self.root.iterdir() + if p.is_file() and (_managed_payload(p) or p.name == 'metadata.sqlite')) + if physical() <= limit: + return + db.execute('VACUUM') + while physical() > limit: + row = db.execute('SELECT id FROM checkpoints ORDER BY used LIMIT 1').fetchone() + with db: + if row: + self._delete(db, row[0]) + self.metrics['evictions'] = self.metrics.get('evictions', 0) + 1 + self._cleanup(db) + else: + db.execute('DELETE FROM tokenizations') + db.execute('DELETE FROM fingerprints') + db.execute('VACUUM') + if not row: + break + + def _refresh(self, db): + version = db.execute('SELECT value FROM revision').fetchone()[0] + if version != self._version: + self.index = Radix() + for key, tokens in db.execute('SELECT id,tokens FROM checkpoints WHERE ns=?', (self.namespace,)): + try: + self.index.insert(json.loads(tokens), key) + except (ValueError, TypeError): + self._delete(db, key) + self._version = version + + def _delete(self, db, key): + row = db.execute('SELECT ns,tokens FROM checkpoints WHERE id=?', (key,)).fetchone() + if row and row[0] == self.namespace: + try: + self.index.remove(json.loads(row[1])) + except (ValueError, TypeError): + pass + db.execute('DELETE FROM checkpoints WHERE id=?', (key,)) + db.execute('DELETE FROM refs WHERE checkpoint=?', (key,)) + db.execute('UPDATE revision SET value=value+1') + + def _cleanup(self, db): + db.execute('DELETE FROM objects WHERE id NOT IN (SELECT object FROM refs)') + live = {r[0] for r in db.execute('SELECT id FROM objects')} + for path in self.root.glob('*.safetensors'): + if _managed_payload(path) and path.stem not in live: + path.unlink(missing_ok=True) + for path in self.root.glob('*.tmp'): + if _managed_payload(path): + path.unlink(missing_ok=True) + + def _evict(self, db): + while self._bytes(db) > self.config.budget_bytes: + row = db.execute('SELECT id FROM checkpoints ORDER BY used LIMIT 1').fetchone() + if row is None: + break + self._delete(db, row[0]) + db.execute('DELETE FROM objects WHERE id NOT IN (SELECT object FROM refs)') + self.metrics['evictions'] = self.metrics.get('evictions', 0) + 1 + self._cleanup(db) + + def _bytes(self, db): + return db.execute('SELECT coalesce(sum(size),0) FROM objects').fetchone()[0] + + def restore(self, tokens, reader): + started = time.perf_counter() + with self.locked() as db: + self._refresh(db) + key = self.index.longest(tokens) + self.metrics['lookup_s'] = time.perf_counter() - started + if key is None: + return None + row = db.execute('SELECT tokens,manifest,checksum FROM checkpoints WHERE id=?', (key,)).fetchone() + try: + if digest(row[0] + row[1]) != row[2]: + raise ValueError('metadata checksum mismatch') + saved, manifest = json.loads(row[0]), json.loads(row[1]) + if saved != list(tokens[:len(saved)]): + raise ValueError('token mismatch') + paths = {} + started = time.perf_counter() + for obj, in db.execute('SELECT object FROM refs WHERE checkpoint=?', (key,)): + path = self.root / (obj + '.safetensors') + if file_hash(path) != obj: + raise ValueError('checksum mismatch') + paths[obj] = path + state = reader(manifest, paths) + db.execute('UPDATE checkpoints SET used=? WHERE id=?', (time.time(), key)) + self.metrics['restore_s'] = time.perf_counter() - started + self.metrics['reused_tokens'] = len(saved) + return saved, state + except (OSError, ValueError, KeyError, TypeError, RuntimeError): + self._delete(db, key) + self._cleanup(db) + self.metrics['corrupt_entries'] = self.metrics.get('corrupt_entries', 0) + 1 + return None + + def publish(self, tokens, writer): + if not tokens: + return + started = time.perf_counter() + key = digest(packed([self.namespace, tokens])) + with self.locked(enforce=True) as db: + self._refresh(db) + if db.execute('SELECT 1 FROM checkpoints WHERE id=?', (key,)).fetchone(): + return + objects = {} + def put(save): + path = self.root / (uuid.uuid4().hex + '.tmp.safetensors') + save(path) + with open(path, 'rb') as f: + os.fsync(f.fileno()) + obj = file_hash(path) + size = path.stat().st_size + dest = self.root / (obj + '.safetensors') + if dest.exists() and file_hash(dest) == obj: + path.unlink() + else: + os.replace(path, dest) + objects[obj] = size + if sum(objects.values()) > self.config.budget_bytes: + raise _Oversized('checkpoint exceeds disk budget') + return obj + try: + manifest = writer(put) + if sum(objects.values()) > self.config.budget_bytes: + self.metrics['oversized'] = self.metrics.get('oversized', 0) + 1 + return + fd = os.open(self.root, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + encoded_tokens, encoded_manifest = packed(tokens), packed(manifest) + db.execute('INSERT INTO checkpoints VALUES (?,?,?,?,?,?)', + (key, self.namespace, encoded_tokens, encoded_manifest, time.time(), + digest(encoded_tokens + encoded_manifest))) + for obj, size in objects.items(): + db.execute('INSERT OR IGNORE INTO objects VALUES (?,?)', (obj, size)) + db.execute('INSERT INTO refs VALUES (?,?)', (key, obj)) + db.execute('UPDATE revision SET value=value+1') + self.index.insert(tokens, key) + self._evict(db) + self._version = db.execute('SELECT value FROM revision').fetchone()[0] + self.metrics['writes'] = self.metrics.get('writes', 0) + 1 + self.metrics['write_s'] = self.metrics.get('write_s', 0) + time.perf_counter() - started + self.metrics['stored_bytes'] = self._bytes(db) + except _Oversized: + self.metrics['oversized'] = self.metrics.get('oversized', 0) + 1 + finally: + self._cleanup(db) + + def tokenize(self, identity, encode): + key = digest(packed([self.namespace, identity])) + started = time.perf_counter() + with self.locked(enforce=True) as db: + row = db.execute('SELECT value,checksum FROM tokenizations WHERE id=?', (key,)).fetchone() + if row and digest(row[0]) != row[1]: + db.execute('DELETE FROM tokenizations WHERE id=?', (key,)) + row = None + if row: + ids = json.loads(row[0]) + db.execute('UPDATE tokenizations SET used=? WHERE id=?', (time.time(), key)) + else: + ids = list(encode()) + value = packed(ids) + if len(value) <= self.config.token_budget_bytes: + db.execute('INSERT OR REPLACE INTO tokenizations VALUES (?,?,?,?)', (key, value, time.time(), digest(value))) + while db.execute('SELECT coalesce(sum(length(value)),0) FROM tokenizations').fetchone()[0] > self.config.token_budget_bytes: + db.execute('DELETE FROM tokenizations WHERE id=(SELECT id FROM tokenizations ORDER BY used LIMIT 1)') + self.metrics['tokenization_s'] = time.perf_counter() - started + return ids + + def inspect(self): + with self.locked() as db: + return dict(checkpoints=db.execute('SELECT count(*) FROM checkpoints').fetchone()[0], + stored_bytes=self._bytes(db), budget_bytes=self.config.budget_bytes, + metadata_bytes=(self.root / 'metadata.sqlite').stat().st_size, + tokenization_bytes=db.execute('SELECT coalesce(sum(length(value)),0) FROM tokenizations').fetchone()[0]) + + def clear(self): + with self.locked() as db: + db.execute('DELETE FROM checkpoints') + db.execute('DELETE FROM refs') + db.execute('DELETE FROM tokenizations') + db.execute('UPDATE revision SET value=value+1') + self._cleanup(db) + + +def file_hash(path): + h = hashlib.sha256() + with open(path, 'rb') as f: + for chunk in iter(lambda: f.read(8 * 1024**2), b''): + h.update(chunk) + return h.hexdigest() + + +def engine_namespace(engine, store): + """Hash artifacts once; reuse hashes only when full stat identity matches.""" + from dataclasses import asdict + from importlib.metadata import version + paths = set(p for p in Path(engine.dir).rglob('*') if p.is_file() + and store.root not in p.resolve().parents + and p.suffix in {'.json', '.jinja', '.safetensors', '.model', '.txt'}) + for value in (engine.cfg.lora, getattr(engine.cfg.prerouter, 'weights_file', '')): + if value: + paths.add(Path(value)) + hashes = [] + with store.locked(enforce=True) as db: + for path in sorted(paths): + path = path.resolve() + st = path.stat() + identity = packed([st.st_dev, st.st_ino, st.st_size, st.st_mtime_ns, st.st_ctime_ns]).decode() + row = db.execute('SELECT identity,hash FROM fingerprints WHERE path=?', (str(path),)).fetchone() + h = row[1] if row and row[0] == identity else file_hash(path) + db.execute('INSERT OR REPLACE INTO fingerprints VALUES (?,?,?)', (str(path), identity, h)) + hashes.append([str(path), h]) + cfg = asdict(engine.cfg) + cfg.pop('conversation_cache', None) + # Include implementation bytes, templates, custom tokenizer identity and all + # family environment switches conservatively (performance switches included). + sources = [(str(p.relative_to(Path(__file__).parents[1])), file_hash(p)) + for p in sorted(Path(__file__).parents[1].rglob('*.py'))] + tokenizer = engine._tok + backend_tokenizer = getattr(tokenizer, 'backend_tokenizer', None) + tokenizer_hash = None + if backend_tokenizer is not None and hasattr(backend_tokenizer, 'to_str'): + tokenizer_hash = digest(backend_tokenizer.to_str().encode()) + elif tokenizer is not None and hasattr(tokenizer, 'get_vocab'): + tokenizer_hash = digest(packed(tokenizer.get_vocab())) + return digest(packed(dict(format=2, tokenizer_hash=tokenizer_hash, artifacts=hashes, config=cfg, interval=store.config.interval, sources=sources, + tokenizer=repr(type(engine._tok)), template=getattr(engine._tok, 'chat_template', None), + versions={p: version(p) for p in ('mlx', 'mlx-lm', 'tokenizers')}, + environment={k: v for k, v in os.environ.items() if k.startswith(('LING_', 'QWEN_', 'EDGE0_', 'MLX_'))}))) diff --git a/src/edge0/engine/base.py b/src/edge0/engine/base.py index c54ef16..1dd78ab 100644 --- a/src/edge0/engine/base.py +++ b/src/edge0/engine/base.py @@ -42,7 +42,19 @@ def __init__(self, model_dir: str, cfg, tokenizer=None): self._last_logits = None self._cap_mlx_cache() t0 = time.time() + self.conversation_cache = None self._build() + self._processed_tokens = [] + self._cache_phase = "ready" + cache_config = getattr(cfg, "conversation_cache", None) + if cache_config and cache_config.directory: + from pathlib import Path + if Path(cache_config.directory).expanduser().resolve() == Path(model_dir).resolve(): + raise ValueError("use a separate directory for conversation checkpoints") + from edge0.conversation import CheckpointStore + from edge0.conversation.store import engine_namespace + self.conversation_cache = CheckpointStore(cache_config, "initializing") + self.conversation_cache.namespace = engine_namespace(self, self.conversation_cache) print(f"[{self.name}] built in {time.time() - t0:.1f}s", flush=True) # ---- family hooks ----------------------------------------------------- @@ -73,20 +85,34 @@ def _lm_logits(self, h: core.array) -> core.array: def prefill(self, token_ids, chunk_size=None, on_progress=None): if chunk_size is None: chunk_size = self.prefill_chunk + if chunk_size < 1: + raise ValueError("chunk size must be positive") total = len(token_ids) done = 0 self._prefill_active = True try: - for start in range(0, total, chunk_size): - chunk = token_ids[start:start + chunk_size] + start = 0 + while start < total: + size = chunk_size + if self.conversation_cache: + interval = self.conversation_cache.config.interval + size = min(size, interval - self.pos % interval) + chunk = token_ids[start:start + size] self._last_logits = self._forward(chunk) self.pos += len(chunk) done += len(chunk) + start += len(chunk) + if self.conversation_cache: + self._processed_tokens.extend(chunk) + if self.pos % self.conversation_cache.config.interval == 0 and done < total: + self._save_checkpoint('prefill') if on_progress is not None: on_progress(done, total) finally: self._prefill_active = False self._prefill_end() + self._cache_phase = "ready" + self._save_checkpoint() return total def next_logits(self) -> core.array: @@ -97,6 +123,11 @@ def step(self, token_id: int) -> core.array: logits = self._forward([token_id]) self.pos += 1 self._step_post(token_id) + if self.conversation_cache: + self._last_logits = logits + self._processed_tokens.append(token_id) + if self.pos % self.conversation_cache.config.interval == 0: + self._save_checkpoint() return logits def generate(self, token_ids, gen_config: GenerationConfig | None = None, @@ -106,11 +137,52 @@ def generate(self, token_ids, gen_config: GenerationConfig | None = None, gen_config = getattr(self.cfg, "gen", GenerationConfig()) if max_new_tokens is None: max_new_tokens = gen_config.max_new_tokens - if len(token_ids) > 1: + suffix = token_ids + if self.conversation_cache and self.pos == 0: + self.conversation_cache.metrics = {} + from edge0.backends.mlx.checkpoint import read + hit = self.conversation_cache.restore(token_ids, read) + if hit: + saved, (caches, logits, state, pos, phase) = hit + try: + if len(caches) != len(self.cache) or pos != len(saved): + raise ValueError("incompatible checkpoint") + self.cache, self._last_logits, self.pos = caches, logits, pos + family_started = time.perf_counter() + self._restore_checkpoint_family(state) + self.conversation_cache.metrics["restore_s"] += time.perf_counter() - family_started + self._processed_tokens = list(saved) + self._cache_phase = phase + suffix = token_ids[len(saved):] + except (ValueError, KeyError, TypeError, RuntimeError): + self.reset() + self.conversation_cache.metrics['reused_tokens'] = 0 + started = time.perf_counter() + if self.conversation_cache: + if len(suffix) == 1 and self.pos == 0: + # Preserve the existing fresh single-token prompt path. + self._last_logits = self._forward(suffix) + self.pos = 1 + self._processed_tokens = list(suffix) + self._save_checkpoint() + elif suffix: + # Even a one-token restored suffix needs the prefill lifecycle. + self.prefill(suffix) + elif self._cache_phase == 'prefill': + self._prefill_end() + self._cache_phase = 'ready' + self.conversation_cache.metrics['prefill_s'] = max(0, time.perf_counter() - started - self.conversation_cache.metrics.get('write_s', 0)) + self.conversation_cache.metrics['remaining_prefill_tokens'] = len(suffix) + elif len(token_ids) > 1: self.prefill(token_ids) elif len(token_ids) == 1: self._last_logits = self._forward(token_ids) self.pos += 1 + self.last_generation_metrics = { + 'prefill_s': (self.conversation_cache.metrics['prefill_s'] + if self.conversation_cache else time.perf_counter() - started)} + decode_started = time.perf_counter() + writes_before_decode = self.conversation_cache.metrics.get('write_s', 0) if self.conversation_cache else 0 out: list[int] = [] history = list(token_ids) logits = self.next_logits() @@ -147,10 +219,28 @@ def generate(self, token_ids, gen_config: GenerationConfig | None = None, logits = self.step(tid) if on_token is not None: on_token(tid) + decode_s = time.perf_counter() - decode_started + if self.conversation_cache: + decode_s -= self.conversation_cache.metrics.get('write_s', 0) - writes_before_decode + self.last_generation_metrics['decode_s'] = max(0, decode_s) + self.last_generation_metrics['decode_tok_s'] = len(out) / decode_s if decode_s > 0 else 0 + self._save_checkpoint() return out + def _save_checkpoint(self, phase="ready"): + if self.conversation_cache and self._processed_tokens: + from edge0.backends.mlx.checkpoint import write + try: + self.conversation_cache.publish(self._processed_tokens, + lambda put: write(self, put, phase)) + except (OSError, ValueError, RuntimeError) as exc: + self.conversation_cache.metrics["last_write_error"] = str(exc) + self.conversation_cache.metrics["write_errors"] = self.conversation_cache.metrics.get("write_errors", 0) + 1 + def reset(self): self._reset_state() + self._processed_tokens = [] + self._cache_phase = "ready" self.pos = 0 self._last_logits = None diff --git a/src/edge0/engine/checkpoint.py b/src/edge0/engine/checkpoint.py new file mode 100644 index 0000000..7bb4df9 --- /dev/null +++ b/src/edge0/engine/checkpoint.py @@ -0,0 +1,50 @@ +"""Explicit family continuation and expert-staging hooks.""" +FIELDS = ('logits', 'logits_prev', 'pred_inds', 'pred_scores', 'oh', 'oh_prev') + + +def capture(engine, family): + st = engine._pg_state + result = {'pg': {k: getattr(st, k) for k in FIELDS} if st else None, + 'experts': {}, 'blocks': {}} + for li, exp in engine._all_stream_layers.items(): + # Finish fills before observing the next decode's staged set. + exp.wait_staged() + staged = exp._staged_state + result['experts'][li] = dict(last=list(exp.last_used), + staged=sorted(staged[2]) if staged else None, + prefill=exp._last_prefill_topk) + if engine._pg_stager: + result['cur_step'] = engine._pg_stager.cur_step + if family == 'ling': + result['pg_cache'] = engine._pg_stager.pg_cache + for owner in st.owners: + block = engine.cfg.moe_spec.block_of(engine.model, owner) + fields = ('prev_topk_oh', 'last_topk') if family == 'ling' else ('prerouter_m_in', 'prerouter_oh') + result['blocks'][owner] = {k: getattr(block, k, None) for k in fields} + if family == 'ling': + result['blocks'][owner]['m_in_cache'] = getattr(engine.cfg.moe_spec.layer_of(engine.model, owner), 'm_in_cache', None) + return result + + +def restore(engine, state, family): + if engine._pg_state: + for key in FIELDS: + setattr(engine._pg_state, key, state['pg'][key]) + if engine._pg_stager: + engine._pg_stager.cur_step = state['cur_step'] + if family == 'ling': + engine._pg_stager.pg_cache = state['pg_cache'] + for owner, values in state['blocks'].items(): + block = engine.cfg.moe_spec.block_of(engine.model, owner) + for key, value in values.items(): + target = engine.cfg.moe_spec.layer_of(engine.model, owner) if key == 'm_in_cache' else block + setattr(target, key, value) + for li, saved in state['experts'].items(): + exp = engine._all_stream_layers[li] + exp.wait_staged() + exp.reset() + exp.last_used = saved['last'] + exp._last_prefill_topk = saved['prefill'] + if saved['staged'] is not None: + exp.stage_experts(saved['staged']) + exp.wait_staged() diff --git a/src/edge0/engine/ling.py b/src/edge0/engine/ling.py index b8a0344..ca4f09d 100644 --- a/src/edge0/engine/ling.py +++ b/src/edge0/engine/ling.py @@ -233,6 +233,19 @@ def _reset_state(self) -> None: if self._pg_stager is not None: self._pg_stager.reset() self._pg_state.reset() + for owner in self._pg_state.owners: + block = self.cfg.moe_spec.block_of(self.model, owner) + block.prev_topk_oh = None + block.last_topk = None + self.cfg.moe_spec.layer_of(self.model, owner).m_in_cache = None + + def _checkpoint_family_state(self): + from edge0.engine.checkpoint import capture + return capture(self, 'ling') + + def _restore_checkpoint_family(self, state): + from edge0.engine.checkpoint import restore + restore(self, state, 'ling') def _lm_logits(self, h: core.array) -> core.array: return self.model.lm_head(h[0, -1]) diff --git a/src/edge0/engine/qwen.py b/src/edge0/engine/qwen.py index 12e0953..1abce73 100644 --- a/src/edge0/engine/qwen.py +++ b/src/edge0/engine/qwen.py @@ -203,5 +203,13 @@ def _reset_state(self) -> None: if self._pg_state is not None: self._pg_state.reset() + def _checkpoint_family_state(self): + from edge0.engine.checkpoint import capture + return capture(self, 'qwen') + + def _restore_checkpoint_family(self, state): + from edge0.engine.checkpoint import restore + restore(self, state, 'qwen') + def _lm_logits(self, h: core.array) -> core.array: return self._lm.lm_head(h[0, -1]) diff --git a/src/edge0/models/base.py b/src/edge0/models/base.py index 09a14b9..fe610af 100644 --- a/src/edge0/models/base.py +++ b/src/edge0/models/base.py @@ -15,6 +15,7 @@ from pathlib import Path from edge0.config import GenerationConfig +from edge0.conversation import CacheConfig from edge0.moe.spec import MoESpec from edge0.prerouter.spec import PrerouterSpec from edge0.streaming.options import LayerOptions @@ -54,6 +55,7 @@ class ModelConfig: lora_r: int = 16 lora_alpha: float = 32.0 gen: GenerationConfig = field(default_factory=GenerationConfig) + conversation_cache: CacheConfig = field(default_factory=CacheConfig) prefill_chunk: int = 2048 hot_window: int = 4 intra_staging: bool = False diff --git a/src/edge0/server/app.py b/src/edge0/server/app.py index 9d94b1b..894d950 100644 --- a/src/edge0/server/app.py +++ b/src/edge0/server/app.py @@ -69,6 +69,7 @@ def _chat_once(server: QueueServer, payload: dict): "finish_reason": "stop", }], "usage": meta["usage"], + **({"edge0_cache": meta["cache"]} if "cache" in meta else {}), } @@ -98,6 +99,7 @@ def produce(): "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": meta["usage"], + **({"edge0_cache": meta["cache"]} if "cache" in meta else {}), }).encode("utf-8")) events.put(b"data: [DONE]\n\n") except Exception as exc: # pragma: no cover - transport-dependent diff --git a/src/edge0/server/chat.py b/src/edge0/server/chat.py index a1502bd..731fa6a 100644 --- a/src/edge0/server/chat.py +++ b/src/edge0/server/chat.py @@ -151,16 +151,35 @@ def run(self, on_token=None) -> tuple[list[int], dict]: # bad pre-routing cross-token state and KV behind, which makes every # later request collapse from its first token. self.engine.reset() - ids = self.prompt_ids() + tokenize_started = time.perf_counter() + cache = getattr(self.engine, 'conversation_cache', None) + if cache: + ids = cache.tokenize(dict(messages=[m.__dict__ for m in self.req.messages], + thinking=self.req.enable_thinking, + default_thinking=getattr(self.engine, 'think', False)), + self.prompt_ids) + tokenization_s = cache.metrics.get('tokenization_s', 0) + else: + ids = self.prompt_ids() + self.engine.last_tokenization_s = time.perf_counter() - tokenize_started gen = self.gen_config() - tokens = self.engine.generate( - ids, gen_config=gen, on_token=on_token) + try: + tokens = self.engine.generate(ids, gen_config=gen, on_token=on_token) + finally: + if cache: + # Completed/cancelled HTTP requests have no active context. + self.engine.reset() usage = { "prompt_tokens": len(ids), "completion_tokens": len(tokens), "total_tokens": len(ids) + len(tokens), } + if cache: + usage['prompt_tokens_details'] = {'cached_tokens': cache.metrics.get('reused_tokens', 0)} + cache.metrics['tokenization_s'] = tokenization_s meta = {"wall_s": round(time.perf_counter() - t0, 3)} + if cache: + meta["cache"] = dict(cache.metrics) return tokens, {"usage": usage, **meta} diff --git a/tests/test_cache_continuation.py b/tests/test_cache_continuation.py new file mode 100644 index 0000000..0af8131 --- /dev/null +++ b/tests/test_cache_continuation.py @@ -0,0 +1,229 @@ +"""Lossless mixed attention/recurrent state and lifecycle coverage.""" +from types import SimpleNamespace +import pytest + +from edge0.backends import core +from edge0.backends.mlx.checkpoint import KVCache, ArraysCache, read +from edge0.conversation import CacheConfig, CheckpointStore +from edge0.engine.base import Edge0Engine +from edge0.config import GenerationConfig + + +class ToyEngine(Edge0Engine): + """State depends on every token and a prerouter-like recurrent scalar.""" + def _build(self): + self._reset_state() + + def _reset_state(self): + self.cache = [KVCache(), ArraysCache(1)] + self.routing = core.array([0.0]) + + def _forward(self, ids, intra_stage=True): + kv, recurrent = self.cache + for tid in ids: + v = core.array([[[[float(tid)]]]]) + kv.update_and_fetch(v, v) + old = recurrent[0] if recurrent[0] is not None else core.array([0.0]) + recurrent[0] = old * 0.5 + tid + self.routing = self.routing * 0.75 + tid + logits = core.arange(8, dtype=core.float32) * (recurrent[0] + self.routing) + core.eval(logits) + return logits + + def _checkpoint_family_state(self): + return {'routing': self.routing} + + def _restore_checkpoint_family(self, state): + self.routing = state['routing'] + + +def engine(tmp_path, namespace='toy'): + eng = ToyEngine('', SimpleNamespace(prefill_chunk=4)) + eng.conversation_cache = CheckpointStore(CacheConfig(str(tmp_path), interval=4), namespace) + return eng + + +@pytest.mark.parametrize('suffix', [[], [5], [5, 6, 7]]) +def test_exact_suffix_restart_logits(tmp_path, suffix): + a = engine(tmp_path) + a.generate([1, 2, 3, 4], max_new_tokens=0) + b = engine(tmp_path) + b.generate([1, 2, 3, 4] + suffix, max_new_tokens=0) + assert b.conversation_cache.metrics['reused_tokens'] == 4 + reference = ToyEngine('', SimpleNamespace(prefill_chunk=4)) + reference.prefill([1, 2, 3, 4] + suffix) + assert core.allclose(b.next_logits(), reference.next_logits()).item() + assert b.pos == 4 + len(suffix) + + +def test_eos_limit_and_callback_cancellation(tmp_path): + eng = engine(tmp_path) + ids = [1, 2, 3] + assert eng.generate(ids, GenerationConfig(eos_ids=(7,))) == [] + assert eng.pos == 3 + eng.reset() + assert eng.generate(ids, max_new_tokens=2) == [7, 7] + assert eng.pos == 5 + eng.reset() + def cancel(token): + raise RuntimeError('cancelled') + with pytest.raises(RuntimeError, match='cancelled'): + eng.generate(ids, max_new_tokens=8, on_token=cancel) + hit = eng.conversation_cache.restore(ids + [7, 7, 7], read) + assert len(hit[0]) <= 5 + assert hit[1][3] == len(hit[0]) + + +def test_mid_prefill_boundary(tmp_path): + eng = engine(tmp_path) + eng.generate(list(range(1, 11)), max_new_tokens=0) + branch = engine(tmp_path) + branch.generate([1, 2, 3, 4, 99], max_new_tokens=0) + assert branch.conversation_cache.metrics['reused_tokens'] == 4 + ref = ToyEngine('', SimpleNamespace(prefill_chunk=4)) + ref.prefill([1, 2, 3, 4, 99]) + assert core.allclose(branch.next_logits(), ref.next_logits()).item() + + +@pytest.mark.slow +@pytest.mark.parametrize('tier,env', [('edge0-8b', 'EDGE0_8B_MODEL'), ('edge0-35b', 'EDGE0_35B_MODEL')]) +def test_real_weight_continuation(tmp_path, tier, env): + import os + from edge0 import AutoEngine + path = os.environ.get(env) + if not path: + pytest.skip(f'{env} unavailable') + eng = AutoEngine.from_pretrained(path, name=tier, + conversation_cache=CacheConfig(str(tmp_path), interval=32)) + try: + ids = eng._tok.encode('Implement a Python function that groups file paths by extension. ' * 8) + eng.generate(ids, max_new_tokens=0) + prompt_logits = eng.next_logits() + tid = int(core.argmax(prompt_logits).item()) + next_logits = eng.step(tid) + eng.reset() + eng.generate(ids, max_new_tokens=0) + assert eng.conversation_cache.metrics.get('reused_tokens') == len(ids) + assert core.allclose(eng.next_logits(), prompt_logits, rtol=1e-5, atol=1e-5).item() + restored_next = eng.step(tid) + assert core.allclose(restored_next, next_logits, rtol=1e-4, atol=1e-4).item() + # A single token appended to a processed-generation checkpoint must + # resume from that exact state, with the normal prefill lifecycle. + eng._save_checkpoint() + eng.prefill([tid]) + suffix_logits = eng.next_logits() + eng.reset() + eng.generate(ids + [tid, tid], max_new_tokens=0) + assert core.allclose(eng.next_logits(), suffix_logits, rtol=1e-4, atol=1e-4).item() + finally: + eng.close() + + +def test_small_qwen_recurrent_backbone(tmp_path): + """Actual Qwen attention + gated-delta layers with random small weights.""" + from edge0.backends.mlx._impl.qwen3_5 import TextModel, TextModelArgs + class SmallQwen(ToyEngine): + def _build(self): + self.model = TextModel(TextModelArgs(hidden_size=32, intermediate_size=64, + num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, + vocab_size=32, linear_num_value_heads=2, linear_num_key_heads=1, + linear_key_head_dim=16, linear_value_head_dim=16, + full_attention_interval=2, head_dim=16)) + self._reset_state() + def _reset_state(self): + self.cache = self.model.make_cache() + def _forward(self, ids, intra_stage=True): + logits = self.model(core.array(ids)[None, :], cache=self.cache)[0, -1] + core.eval(logits) + return logits + def _checkpoint_family_state(self): + return {} + def _restore_checkpoint_family(self, state): + pass + model = SmallQwen('', SimpleNamespace(prefill_chunk=4)) + model.conversation_cache = CheckpointStore(CacheConfig(str(tmp_path), interval=4), 'small-qwen') + model.generate([1, 2, 3, 4, 5], max_new_tokens=0) + expected = model.step(6) + model.reset() + model.generate([1, 2, 3, 4, 5], max_new_tokens=0) + actual = model.step(6) + assert core.allclose(expected, actual, rtol=1e-5, atol=1e-5).item() + + +@pytest.mark.parametrize('family', ['ling', 'qwen']) +def test_family_prerouter_fields_and_staging(tmp_path, family): + from edge0.engine.checkpoint import capture, restore, FIELDS + from edge0.prerouter.state import PrerouterState + from edge0.backends.mlx.checkpoint import write + block = SimpleNamespace(prev_topk_oh=core.array([2.0]), last_topk=core.array([1]), + prerouter_m_in=core.array([3.0]), prerouter_oh=core.array([4.0])) + layer = SimpleNamespace(m_in_cache=core.array([5.0])) + spec = SimpleNamespace(block_of=lambda model, owner: block, + layer_of=lambda model, owner: layer) + class Expert: + last_used = [1, 2] + _staged_state = ((), [], {2, 3}) + _last_prefill_topk = core.array([1, 2]) + def wait_staged(self): + pass + def reset(self): + self._staged_state = None + def stage_experts(self, ids): + self._staged_state = ((), [], set(ids)) + eng = engine(tmp_path) + eng.generate([1, 2], max_new_tokens=0) + eng.cfg.moe_spec = spec + eng.model = None + eng._pg_state = PrerouterState(2, 4, 1, owners=[0]) + for i, key in enumerate(FIELDS): + setattr(eng._pg_state, key, [core.array([float(i)]), None]) + eng._pg_stager = SimpleNamespace(cur_step=7, pg_cache={1: core.array([8.0])}) + eng._all_stream_layers = {1: Expert()} + eng._checkpoint_family_state = lambda: capture(eng, family) + eng.conversation_cache.clear() + eng.conversation_cache.publish([1, 2], lambda put: write(eng, put)) + state = eng.conversation_cache.restore([1, 2], read)[1][2] + eng._pg_state.reset() + eng._pg_stager.pg_cache = {} + restore(eng, state, family) + for i, key in enumerate(FIELDS): + assert getattr(eng._pg_state, key)[0].item() == i + assert eng._pg_stager.cur_step == 7 + assert eng._all_stream_layers[1]._staged_state[2] == {2, 3} + if family == 'ling': + assert eng._pg_stager.pg_cache[1].item() == 8 + + +def test_chat_releases_idle_state_and_reports_usage(tmp_path): + from edge0.server.chat import ChatSession, ChatRequest, ChatMessage + eng = engine(tmp_path) + eng._tok = SimpleNamespace(encode=lambda text: [1, 2, 3], bos_token_id=0) + request = ChatRequest('toy', [ChatMessage('user', 'hello')], max_tokens=2) + _, first = ChatSession(eng, request).run() + assert eng.pos == 0 + assert all(c.empty() for c in eng.cache) + _, repeated = ChatSession(eng, request).run() + assert repeated['usage']['prompt_tokens'] == 3 + assert repeated['usage']['total_tokens'] == 5 + assert repeated['usage']['prompt_tokens_details']['cached_tokens'] == 3 + assert eng.pos == 0 + def cancelled(token): + raise RuntimeError('cancelled') + with pytest.raises(RuntimeError): + ChatSession(eng, request).run(on_token=cancelled) + assert eng.pos == 0 + assert all(c.empty() for c in eng.cache) + + +def test_sampling_penalty_history_keeps_reused_prompt(tmp_path, monkeypatch): + eng = engine(tmp_path) + eng.generate([1, 2, 3, 4], max_new_tokens=0) + eng.reset() + seen = [] + def sample(logits, **kwargs): + seen.append(list(kwargs['history'])) + return 6 + monkeypatch.setattr('edge0.engine.base.sample', sample) + eng.generate([1, 2, 3, 4, 5], GenerationConfig( + first_token_greedy=False, repetition_penalty=1.1, max_new_tokens=2)) + assert seen == [[1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6]] diff --git a/tests/test_conversation_cache.py b/tests/test_conversation_cache.py new file mode 100644 index 0000000..cbe6856 --- /dev/null +++ b/tests/test_conversation_cache.py @@ -0,0 +1,210 @@ +"""Persistence/index tests run without loading model weights.""" +import json +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from edge0.conversation import CacheConfig, CheckpointStore +from edge0.conversation.store import Radix + + +def writer(*payloads): + def write(put): + return [put(lambda path, data=data: path.write_bytes(data)) for data in payloads] + return write + + +def reader(manifest, paths): + return [paths[key].read_bytes() for key in manifest] + + +def test_radix_branches_and_exact(): + tree = Radix() + for tokens in ([1, 2, 3], [1, 2, 4], [1], [1, 2, 3, 5], [8, 9]): + tree.insert(tokens, tokens) + assert tree.longest([1, 2]) == [1] + assert tree.longest([1, 2, 3]) == [1, 2, 3] + assert tree.longest([1, 2, 3, 5, 6]) == [1, 2, 3, 5] + assert tree.longest([8]) is None + assert tree.longest([2]) is None + + +def test_restart_namespace_corruption_and_orphans(tmp_path): + cfg = CacheConfig(str(tmp_path)) + store = CheckpointStore(cfg, 'a') + store.publish([1, 2], writer(b'first')) + (tmp_path / ('a' * 32 + '.tmp')).write_bytes(b'partial') + (tmp_path / ('f' * 64 + '.safetensors')).write_bytes(b'orphan') + store = CheckpointStore(cfg, 'a') + assert not list(tmp_path.glob('*.tmp')) + assert store.restore([1, 2, 3], reader) == ([1, 2], [b'first']) + assert CheckpointStore(cfg, 'b').restore([1, 2], reader) is None + next(tmp_path.glob('*.safetensors')).write_bytes(b'corrupt') + assert store.restore([1, 2], reader) is None + assert store.inspect()['checkpoints'] == 0 + + +def test_shared_eviction_and_oversized(tmp_path): + store = CheckpointStore(CacheConfig(str(tmp_path), budget_bytes=10), 'a') + store.publish([1], writer(b'abc', b'def')) + store.publish([2], writer(b'abc', b'ghi')) + assert store.inspect()['stored_bytes'] == 9 + store.publish([3], writer(b'abc', b'jkl')) + assert store.restore([1], reader) is None + assert store.restore([2], reader)[1] == [b'abc', b'ghi'] + store.publish([4], writer(b'x' * 11)) + assert store.restore([4], reader) is None + assert store.inspect()['stored_bytes'] <= 10 + assert sum(p.stat().st_size for p in tmp_path.glob('*.safetensors')) <= 10 + + +def test_failed_writer_never_published(tmp_path): + store = CheckpointStore(CacheConfig(str(tmp_path)), 'a') + def fail(put): + put(lambda p: p.write_bytes(b'first')) + raise RuntimeError('interrupted') + with pytest.raises(RuntimeError): + store.publish([1], fail) + assert store.restore([1], reader) is None + assert not list(tmp_path.glob('*.safetensors')) + + +def test_concurrent_instances_refresh(tmp_path): + cfg = CacheConfig(str(tmp_path)) + stores = [CheckpointStore(cfg, 'a') for _ in range(4)] + def worker(i): + stores[i % 4].publish([i, i + 1], writer(str(i).encode())) + with ThreadPoolExecutor(4) as pool: + list(pool.map(worker, range(40))) + assert stores[0].inspect()['checkpoints'] == 40 + for i in range(40): + assert stores[0].restore([i, i + 1, 99], reader)[1] == [str(i).encode()] + + +def test_tokenization_bounded_and_exact(tmp_path): + store = CheckpointStore(CacheConfig(str(tmp_path), token_budget_bytes=10), 'a') + calls = [] + def encode(): + calls.append(1) + return [1, 2] + assert store.tokenize({'thinking': False}, encode) == [1, 2] + store.tokenize({'thinking': False}, encode) + assert len(calls) == 1 + store.tokenize({'thinking': True}, encode) + store.tokenize({'messages': ['different']}, encode) + assert len(calls) == 3 + assert store.inspect()['tokenization_bytes'] <= 10 + + +def test_artifact_identity_validation(tmp_path, monkeypatch): + from dataclasses import dataclass + from types import SimpleNamespace + from edge0.conversation import store as module + @dataclass + class Config: + lora: str = '' + prerouter: object = None + model = tmp_path / 'model' + model.mkdir() + artifact = model / 'model.safetensors' + artifact.write_bytes(b'weights') + engine = SimpleNamespace(dir=str(model), cfg=Config(), _tok=None) + store = CheckpointStore(CacheConfig(str(tmp_path / 'cache')), '') + calls = [] + original = module.file_hash + def track(path): + if path == artifact: + calls.append(path) + return original(path) + monkeypatch.setattr(module, 'file_hash', track) + first = module.engine_namespace(engine, store) + assert module.engine_namespace(engine, store) == first + assert len(calls) == 1 + artifact.write_bytes(b'changed') + assert module.engine_namespace(engine, store) != first + assert len(calls) == 2 + + +def test_missing_payload(tmp_path): + store = CheckpointStore(CacheConfig(str(tmp_path)), 'a') + store.publish([1], writer(b'one')) + next(tmp_path.glob('*.safetensors')).unlink() + assert store.restore([1], reader) is None + + +def test_multiprocess_access(tmp_path): + import subprocess + import sys + program = ''' +import sys +from edge0.conversation import CacheConfig, CheckpointStore +s = CheckpointStore(CacheConfig(sys.argv[1]), 'a') +for i in range(12): + s.publish([int(sys.argv[2]), i], lambda put: [put(lambda p: p.write_bytes(b'shared'))]) +''' + workers = [subprocess.Popen([sys.executable, '-c', program, str(tmp_path), str(i)]) for i in range(3)] + for worker in workers: + assert worker.wait(timeout=30) == 0 + store = CheckpointStore(CacheConfig(str(tmp_path)), 'a') + assert store.inspect()['checkpoints'] == 36 + assert store.inspect()['stored_bytes'] == len(b'shared') + + +def test_metadata_corruption_and_shared_repair(tmp_path): + store = CheckpointStore(CacheConfig(str(tmp_path)), 'a') + store.publish([1], writer(b'shared')) + store.publish([2], writer(b'shared')) + next(tmp_path.glob('*.safetensors')).write_bytes(b'broken') + assert store.restore([1], reader) is None + store.publish([1], writer(b'shared')) + assert store.restore([2], reader)[1] == [b'shared'] + with store.locked() as db: + db.execute("UPDATE checkpoints SET manifest=?", (b'[]',)) + assert store.restore([1], reader) is None + + +def test_cleanup_preserves_unmanaged_files(tmp_path): + (tmp_path / 'model.safetensors').write_bytes(b'not a cache payload') + (tmp_path / 'user.tmp').write_bytes(b'not a cache temporary') + store = CheckpointStore(CacheConfig(str(tmp_path)), 'a') + store.clear() + assert (tmp_path / 'model.safetensors').exists() + assert (tmp_path / 'user.tmp').exists() + + +def test_local_writes_update_index_without_rebuild(tmp_path, monkeypatch): + store = CheckpointStore(CacheConfig(str(tmp_path)), 'a') + store.publish([1], writer(b'one')) + index = store.index + store.publish([1, 2], writer(b'two')) + assert store.restore([1, 2, 3], reader)[0] == [1, 2] + assert store.index is index + + +def test_physical_budget_with_metadata(tmp_path): + cfg = CacheConfig(str(tmp_path), budget_bytes=2048, token_budget_bytes=1024) + store = CheckpointStore(cfg, 'a') + for i in range(20): + store.publish([i] * 300, writer(bytes([i]) * 1000)) + store.tokenize(str(i), lambda: list(range(100))) + physical = sum(p.stat().st_size for p in tmp_path.iterdir() if p.is_file()) + assert physical <= cfg.budget_bytes + cfg.token_budget_bytes + 65536 + + +def test_process_exit_before_metadata_publication(tmp_path): + import subprocess + import sys + program = ''' +import os, sys +from edge0.conversation import CacheConfig, CheckpointStore +s = CheckpointStore(CacheConfig(sys.argv[1]), 'a') +def write(put): + put(lambda p: p.write_bytes(b'unpublished')) + os._exit(9) +s.publish([1, 2], write) +''' + result = subprocess.run([sys.executable, '-c', program, str(tmp_path)], timeout=30) + assert result.returncode == 9 + store = CheckpointStore(CacheConfig(str(tmp_path)), 'a') + assert store.restore([1, 2], reader) is None + assert not list(tmp_path.glob('*.safetensors'))