lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/tools/cache_simulator/gen_bench_dataset.py

@@ -0,0 +1,360 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Convert LMCache lookup-hash JSONL logs into a ``vllm bench serve`` custom
+dataset (JSONL with ``"prompt"`` and ``"output_tokens"`` fields).
+
+The conversion preserves the **prefix-sharing structure** of the original
+requests: requests that shared a chunk hash in the lookup logs will share the
+same token prefix in the generated prompts.  This allows the synthetic
+benchmark to exercise LMCache prefix caching in the same pattern as the
+original production workload.
+
+Algorithm
+---------
+1. Build a *safe vocabulary*: token IDs from the tokenizer whose single-token
+   round-trip is stable (``decode([id])`` re-encodes back to exactly ``[id]``).
+   Tokens that start with a leading space are preferred to avoid BPE merges at
+   chunk boundaries.
+2. For each unique chunk hash, deterministically seed an RNG from
+   ``SHA-256(hash)`` and sample ``chunk_size`` token IDs from the safe vocab.
+   The same hash always produces the same token sequence.
+3. Per request: concatenate the token sequences for each full chunk, add
+   ``tail_len = seq_len mod chunk_size`` tail tokens (unique per request so
+   they are never accidentally cached), decode the whole list to text, and
+   write ``{"prompt": <text>, "output_tokens": <N>}``.
+
+Usage (module mode)::
+
+    python3 -m lmcache.tools.cache_simulator.gen_bench_dataset \\
+        -i /path/to/lookup_hashes/ \\
+        --tokenizer /models/DeepSeek-V3 \\
+        --output-len 128 \\
+        -o bench_dataset.jsonl
+
+Usage (CLI)::
+
+    lmcache tool cache-simulator gen-dataset \\
+        -i /path/to/lookup_hashes/ \\
+        --tokenizer /models/DeepSeek-V3 \\
+        --output-len 128 \\
+        -o bench_dataset.jsonl
+"""
+
+# Standard
+from pathlib import Path
+from typing import Any
+import argparse
+import hashlib
+import json
+import random
+import sys
+
+# ---------------------------------------------------------------------------
+# Safe-vocabulary helpers
+# ---------------------------------------------------------------------------
+
+
+def build_safe_vocab(tokenizer: Any) -> list[int]:
+    """Return token IDs that round-trip stably through the tokenizer.
+
+    A *safe* token is one where::
+
+        tokenizer.encode(tokenizer.decode([id]), add_special_tokens=False) == [id]
+
+    These tokens can be concatenated safely — decoding a list of safe tokens
+    and re-tokenizing it will reproduce the same token IDs as long as each
+    token starts with a leading space (preventing BPE merges at boundaries).
+
+    Tokens that begin with a space (``Ġ``, ``▁``, or ``" "``) are placed
+    first so that ``hash_to_tokens`` preferentially picks space-prefixed
+    tokens, making BPE boundary merges much less likely.
+
+    Args:
+        tokenizer: A HuggingFace ``PreTrainedTokenizer`` or
+            ``PreTrainedTokenizerFast`` instance.
+
+    Returns:
+        A non-empty list of safe token IDs.  Raises ``RuntimeError`` if
+        fewer than 10 safe tokens are found (the tokenizer is probably
+        unsupported).
+    """
+    vocab_size = tokenizer.vocab_size
+    safe_with_space: list[int] = []
+    safe_other: list[int] = []
+
+    for token_id in range(vocab_size):
+        try:
+            decoded = tokenizer.decode(
+                [token_id], skip_special_tokens=True, clean_up_tokenization_spaces=False
+            )
+        except Exception:
+            continue
+        if not decoded or not decoded.isprintable():
+            continue
+        # Re-encode and check round-trip
+        try:
+            re_encoded = tokenizer.encode(decoded, add_special_tokens=False)
+        except Exception:
+            continue
+        if re_encoded != [token_id]:
+            continue
+        # Prefer tokens that start with whitespace (prevents inter-chunk merges)
+        if decoded[0] in (" ", "\u0120", "\u2581"):
+            safe_with_space.append(token_id)
+        else:
+            safe_other.append(token_id)
+
+    result = safe_with_space + safe_other
+    if len(result) < 10:
+        raise RuntimeError(
+            f"Only {len(result)} safe tokens found in the vocabulary. "
+            "This tokenizer may not be supported."
+        )
+    return result
+
+
+def hash_to_tokens(chunk_hash: str, n_tokens: int, vocab_ids: list[int]) -> list[int]:
+    """Deterministically map *chunk_hash* to *n_tokens* token IDs.
+
+    The mapping is a pure function of ``chunk_hash`` and ``n_tokens`` — the
+    same inputs always produce the same output regardless of call order.
+
+    Args:
+        chunk_hash: Hex string identifying the chunk (e.g. ``"0xabcd1234"``).
+        n_tokens: Number of token IDs to generate.
+        vocab_ids: Pool of safe token IDs to sample from.
+
+    Returns:
+        A list of ``n_tokens`` token IDs drawn from ``vocab_ids``.
+    """
+    seed_bytes = hashlib.sha256(chunk_hash.encode()).digest()
+    seed = int.from_bytes(seed_bytes, "big") % (2**31)
+    rng = random.Random(seed)
+    return [rng.choice(vocab_ids) for _ in range(n_tokens)]
+
+
+# ---------------------------------------------------------------------------
+# Main conversion
+# ---------------------------------------------------------------------------
+
+
+def gen_bench_dataset(
+    events: list[dict[str, Any]],
+    tokenizer: Any,
+    output_len: int,
+    output_path: Path,
+) -> int:
+    """Convert lookup-hash events into a vllm bench serve custom dataset.
+
+    Each event becomes one JSONL line with fields ``"prompt"`` (text) and
+    ``"output_tokens"`` (int).  Chunk hashes are deterministically mapped to
+    token sequences so that prefix-sharing in the original workload is
+    reproduced in the synthetic prompts.
+
+    Args:
+        events: Lookup-hash events as returned by
+            :func:`~lmcache.tools.cache_simulator.simulator.load_lookup_events`.
+        tokenizer: A HuggingFace tokenizer loaded for the target model.
+        output_len: Number of output tokens for every request in the dataset.
+        output_path: Where to write the output JSONL file.
+
+    Returns:
+        Number of records written.
+    """
+    print("Building safe vocabulary …", flush=True)
+    vocab_ids = build_safe_vocab(tokenizer)
+    space_chars = (" ", "\u0120", "\u2581")
+    n_space = sum(1 for t in vocab_ids if tokenizer.decode([t])[0] in space_chars)
+    print(
+        f"  Safe vocab size: {len(vocab_ids)} tokens ({n_space} space-prefixed)",
+        flush=True,
+    )
+
+    # Cache: chunk_hash -> token_id list (memoized across requests)
+    hash_cache: dict[str, list[int]] = {}
+    n_written = 0
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "w", encoding="utf-8") as fout:
+        for event in events:
+            chunk_hashes: list[str] = event.get("chunk_hashes", [])
+            chunk_size: int = event["chunk_size"]
+            seq_len: int = event["seq_len"]
+
+            # --- Full-chunk tokens ----------------------------------------
+            all_token_ids: list[int] = []
+            for h in chunk_hashes:
+                if h not in hash_cache:
+                    hash_cache[h] = hash_to_tokens(h, chunk_size, vocab_ids)
+                all_token_ids.extend(hash_cache[h])
+
+            # --- Tail tokens ----------------------------------------------
+            # seq_len mod chunk_size tail tokens; always a miss in LMCache.
+            # Use a per-request unique seed so tails are never accidentally
+            # shared across requests.
+            tail_len = seq_len - len(chunk_hashes) * chunk_size
+            if tail_len > 0:
+                # Build a per-request identifier: prefer request_id, fall back
+                # to timestamp+index to ensure uniqueness.
+                req_id = event.get(
+                    "request_id", f"{event.get('timestamp', n_written)}_{n_written}"
+                )
+                tail_hash = f"__tail__{req_id}"
+                all_token_ids.extend(hash_to_tokens(tail_hash, tail_len, vocab_ids))
+
+            # --- Decode and write -----------------------------------------
+            prompt_text = tokenizer.decode(
+                all_token_ids,
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=False,
+            )
+            record = {"prompt": prompt_text, "output_tokens": output_len}
+            fout.write(json.dumps(record, ensure_ascii=False) + "\n")
+            n_written += 1
+
+            if n_written % 1000 == 0:
+                print(f"  {n_written} / {len(events)} records written …", flush=True)
+
+    return n_written
+
+
+# ---------------------------------------------------------------------------
+# CLI helpers (shared by main() and lmcache tool cache-simulator gen-dataset)
+# ---------------------------------------------------------------------------
+
+
+def add_gen_dataset_arguments(parser: argparse.ArgumentParser) -> None:
+    """Register all gen-dataset CLI flags onto *parser*.
+
+    Called by both :func:`main` (``python -m`` entry point) and the
+    ``lmcache tool cache-simulator gen-dataset`` sub-command in
+    ``lmcache/cli/commands/tool/cache_simulator.py``.
+
+    **When adding or removing a flag**, edit only this function — the
+    ``lmcache tool`` command picks up the change automatically.
+
+    Args:
+        parser: The :class:`argparse.ArgumentParser` (or sub-parser) to
+            populate.
+    """
+    parser.add_argument(
+        "-i",
+        "--input",
+        nargs="+",
+        required=True,
+        metavar="PATH",
+        type=Path,
+        help=(
+            "One or more JSONL files or directories containing "
+            "``lookup_hashes_*.jsonl`` files."
+        ),
+    )
+    parser.add_argument(
+        "--tokenizer",
+        required=True,
+        metavar="PATH",
+        help=(
+            "Path or HuggingFace model name for the tokenizer used to decode "
+            "synthetic token IDs into text.  Should match the model that will "
+            "be benchmarked."
+        ),
+    )
+    parser.add_argument(
+        "--output-len",
+        type=int,
+        default=128,
+        metavar="N",
+        help="Number of output tokens for every request (default: %(default)s).",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        default="bench_dataset.jsonl",
+        metavar="FILE",
+        help="Output JSONL file path (default: %(default)s).",
+    )
+    parser.add_argument(
+        "-n",
+        "--max-samples",
+        type=int,
+        default=None,
+        metavar="N",
+        help="Truncate to the first N events after sorting by timestamp.",
+    )
+    parser.add_argument(
+        "--model",
+        default=None,
+        metavar="NAME",
+        help=(
+            "Filter events by ``model_name`` (exact match).  Useful when logs "
+            "contain traffic for multiple models."
+        ),
+    )
+
+
+def run_gen_dataset(args: argparse.Namespace) -> None:
+    """Execute the gen-dataset workflow from a parsed argument namespace.
+
+    Loads events via
+    :func:`~lmcache.tools.cache_simulator.simulator.load_lookup_events`,
+    loads the tokenizer, calls :func:`gen_bench_dataset`, and prints a
+    summary.
+
+    Args:
+        args: Parsed CLI namespace produced by a parser that was populated
+            with :func:`add_gen_dataset_arguments`.
+    """
+    # Lazy imports — keeps CLI startup fast
+    # Third Party
+    from transformers import AutoTokenizer
+
+    # First Party
+    from lmcache.tools.cache_simulator.simulator import load_lookup_events
+
+    # ---- Load events -------------------------------------------------------
+    print(f"Loading events from: {[str(p) for p in args.input]}", flush=True)
+    events = load_lookup_events(
+        args.input,
+        model=args.model,
+        max_samples=args.max_samples,
+    )
+    if not events:
+        print("No events found — nothing to write.", file=sys.stderr)
+        sys.exit(1)
+    print(f"Loaded {len(events):,} events.", flush=True)
+
+    # ---- Load tokenizer ----------------------------------------------------
+    print(f"Loading tokenizer: {args.tokenizer}", flush=True)
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
+
+    # ---- Convert -----------------------------------------------------------
+    output_path = Path(args.output)
+    n = gen_bench_dataset(
+        events=events,
+        tokenizer=tokenizer,
+        output_len=args.output_len,
+        output_path=output_path,
+    )
+    print(f"\nWrote {n:,} records to {output_path}")
+
+
+# ---------------------------------------------------------------------------
+# Module entry point
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    """CLI entry point for ``python -m``."""
+    parser = argparse.ArgumentParser(
+        description=(
+            "Convert LMCache lookup-hash JSONL logs into a "
+            "vllm bench serve custom dataset."
+        )
+    )
+    add_gen_dataset_arguments(parser)
+    args = parser.parse_args()
+    run_gen_dataset(args)
+
+
+if __name__ == "__main__":
+    main()

lmcache/tools/cache_simulator/lru_cache.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LRU cache implementations for the cache simulator.
+
+Two variants are provided:
+
+* ``LRUCacheFast`` — O(1) OrderedDict-backed.  Supports only hit/miss queries.
+  Use this for capacity sweeps where per-chunk statistics are not needed.
+
+* ``LRUCache`` — O(log n) dict + SortedList.  Adds ``position(key)`` for
+  computing cache-position statistics (0 = MRU, capacity-1 = LRU).
+  Use this for the detailed single-run report.
+"""
+
+# Standard
+from collections import OrderedDict
+from typing import Any
+
+# Third Party
+from sortedcontainers import SortedList
+
+
+class LRUCacheFast:
+    """
+    Lightweight LRU cache using a single :class:`OrderedDict`.
+
+    All operations are O(1).  No per-key position tracking.
+
+    Parameters
+    ----------
+    capacity:
+        Maximum number of entries before LRU eviction kicks in.
+    """
+
+    def __init__(self, capacity: int) -> None:
+        if capacity < 1:
+            raise ValueError(f"LRUCacheFast capacity must be >= 1, got {capacity}")
+        self.capacity = capacity
+        self._cache: OrderedDict[Any, None] = OrderedDict()
+        self.eviction_count: int = 0
+
+    def contains(self, key: Any) -> bool:
+        return key in self._cache
+
+    def access(self, key: Any) -> None:
+        """Mark an existing entry as most-recently used. O(1)."""
+        self._cache.move_to_end(key)
+
+    def insert(self, key: Any) -> None:
+        """Insert a new entry, or refresh if already present. O(1)."""
+        if key in self._cache:
+            self._cache.move_to_end(key)
+            return
+        if len(self._cache) >= self.capacity:
+            self._cache.popitem(last=False)
+            self.eviction_count += 1
+        self._cache[key] = None
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+
+class LRUCache:
+    """
+    LRU cache backed by a dict (O(1) lookup) and a
+    :class:`~sortedcontainers.SortedList` (O(log n) rank queries).
+
+    Each entry carries a strictly-increasing clock value so that the SortedList
+    order is unambiguous.
+
+    * LRU end = smallest clock value = ``SortedList[0]``
+    * MRU end = largest clock value  = ``SortedList[-1]``
+
+    Parameters
+    ----------
+    capacity:
+        Maximum number of entries before LRU eviction kicks in.
+    """
+
+    def __init__(self, capacity: int) -> None:
+        if capacity < 1:
+            raise ValueError(f"LRUCache capacity must be >= 1, got {capacity}")
+        self.capacity = capacity
+        self._clock: int = 0
+        self._map: dict[Any, int] = {}  # key -> clock value at last access
+        self._sl: SortedList = SortedList(key=lambda x: x[0])
+        self.eviction_count: int = 0
+
+    def contains(self, key: Any) -> bool:
+        return key in self._map
+
+    def position(self, key: Any) -> int:
+        """
+        LRU rank of *key*: 0 = most-recently used, ``len-1`` = least-recently used.
+        O(log n).
+        """
+        clock = self._map[key]
+        idx = self._sl.index((clock, key))
+        return len(self._sl) - 1 - idx
+
+    def access(self, key: Any) -> None:
+        """Mark an existing entry as most-recently used. O(log n)."""
+        old_clock = self._map[key]
+        self._sl.remove((old_clock, key))
+        self._clock += 1
+        self._map[key] = self._clock
+        self._sl.add((self._clock, key))
+
+    def insert(self, key: Any) -> None:
+        """Insert a new entry, or refresh if already present. O(log n)."""
+        if key in self._map:
+            self.access(key)
+            return
+        if len(self._map) >= self.capacity:
+            lru_clock, lru_key = self._sl[0]
+            self._sl.remove((lru_clock, lru_key))
+            del self._map[lru_key]
+            self.eviction_count += 1
+        self._clock += 1
+        self._map[key] = self._clock
+        self._sl.add((self._clock, key))
+
+    def __len__(self) -> int:
+        return len(self._map)