@fugood/buttress-server 2.26.0-beta.3 → 2.26.0-beta.5

package/lib/mlx-bridge.py

@@ -0,0 +1,681 @@
+#!/usr/bin/env python3
+"""MLX-VLM IPC bridge for buttress-backend-core.
+
+Protocol: newline-delimited JSON over stdin/stdout.
+Request:  {"id": "...", "method": "...", "params": {...}}
+Response: {"id": "...", "result": {...}} or {"id": "...", "error": {"message": "..."}}
+Stream:   {"id": "...", "event": "token"|"progress"|"result", "data": {...}}
+
+Logging goes to stderr to avoid mixing with IPC.
+"""
+
+import json
+import sys
+import os
+import hashlib
+import time
+import threading
+import traceback
+
+# Shared state
+_model = None
+_processor = None
+_config = None
+_model_id = None
+_is_vlm = False  # True when loaded via mlx-vlm, False when loaded via mlx-lm
+_session_cache = None
+_cancel_flags = set()
+_cancel_lock = threading.Lock()
+_output_lock = threading.Lock()
+
+
+# ---------------------------------------------------------------------------
+# Session Cache Manager
+# ---------------------------------------------------------------------------
+
+class SessionCacheManager:
+    """Manages KV cache persistence for prompt prefix reuse.
+
+    Saves prompt-only KV state to .safetensors files after each generation.
+    On subsequent requests, loads the longest matching prefix cache so the
+    model only needs to process new tokens.
+    """
+
+    def __init__(self, cache_dir, max_entries=100, max_size_bytes=5 * 1024 * 1024 * 1024):
+        self.cache_dir = cache_dir
+        self.max_entries = max_entries
+        self.max_size_bytes = max_size_bytes
+        self.map_path = os.path.join(cache_dir, "cache-map.json")
+        self.states_dir = os.path.join(cache_dir, "states")
+        self.cache_map = None
+        # Keep the last-used cache in memory to avoid re-reading from disk
+        self._mem_cache = None
+        self._mem_cache_key = None
+
+    def initialize(self):
+        os.makedirs(self.states_dir, exist_ok=True)
+        self._load_map()
+        count = len(self.cache_map.get("entries", {}))
+        log(f"Session cache initialized: {count} entries in {self.cache_dir}")
+
+    # -- Map persistence --
+
+    def _load_map(self):
+        try:
+            with open(self.map_path, "r") as f:
+                self.cache_map = json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            self.cache_map = {"entries": {}, "totalSize": 0}
+
+    def _save_map(self):
+        try:
+            with open(self.map_path, "w") as f:
+                json.dump(self.cache_map, f)
+        except Exception as e:
+            log(f"Failed to save cache map: {e}")
+
+    @staticmethod
+    def _entry_id(text):
+        return hashlib.sha256(text.encode()).hexdigest()[:32]
+
+    # -- Lookup --
+
+    def find_best_prefix(self, prompt_text):
+        """Return the entry whose promptText is the longest prefix of *prompt_text*."""
+        best = None
+        best_len = 0
+        for entry in self.cache_map.get("entries", {}).values():
+            saved = entry.get("promptText", "")
+            if prompt_text.startswith(saved) and len(saved) > best_len:
+                best = entry
+                best_len = len(saved)
+        return best
+
+    def load_cache_for(self, entry):
+        """Load a saved prompt cache from disk (or memory)."""
+        from mlx_lm.models.cache import load_prompt_cache
+
+        entry_id = entry["id"]
+
+        # Fast path: still in memory
+        if self._mem_cache is not None and self._mem_cache_key == entry_id:
+            entry["lastAccessedAt"] = time.time()
+            self._save_map()
+            log(f"Session cache hit (memory): {entry_id}")
+            return self._mem_cache
+
+        file_path = entry.get("filePath")
+        if not file_path or not os.path.exists(file_path):
+            self._remove_entry(entry_id)
+            return None
+
+        try:
+            cache = load_prompt_cache(file_path)
+            entry["lastAccessedAt"] = time.time()
+            self._save_map()
+            self._mem_cache = cache
+            self._mem_cache_key = entry_id
+            return cache
+        except Exception as e:
+            log(f"Failed to load cache {file_path}: {e}")
+            self._remove_entry(entry_id)
+            return None
+
+    # -- Save --
+
+    def save_after_generation(self, prompt_text, prompt_tokens, cache, generated_tokens):
+        """Trim generated tokens from *cache* and persist the prompt-only KV state."""
+        from mlx_lm.cache_prompt import save_prompt_cache
+        from mlx_lm.models.cache import trim_prompt_cache
+
+        if generated_tokens > 0:
+            trim_prompt_cache(cache, generated_tokens)
+
+        entry_id = self._entry_id(prompt_text)
+        file_path = os.path.join(self.states_dir, f"{entry_id}.safetensors")
+
+        try:
+            save_prompt_cache(file_path, cache, metadata={
+                "prompt_tokens": str(prompt_tokens),
+                "prompt_length": str(len(prompt_text)),
+            })
+            file_size = os.path.getsize(file_path)
+
+            # Remove old entry size from total
+            old = self.cache_map["entries"].get(entry_id)
+            if old:
+                self.cache_map["totalSize"] -= old.get("fileSize", 0)
+
+            self.cache_map["entries"][entry_id] = {
+                "id": entry_id,
+                "promptText": prompt_text,
+                "promptTokens": prompt_tokens,
+                "filePath": file_path,
+                "fileSize": file_size,
+                "createdAt": old["createdAt"] if old else time.time(),
+                "lastAccessedAt": time.time(),
+            }
+            self.cache_map["totalSize"] += file_size
+
+            # Keep in memory
+            self._mem_cache = cache
+            self._mem_cache_key = entry_id
+
+            self._evict_superseded(prompt_text, entry_id)
+            self._evict_if_needed()
+            self._save_map()
+
+            log(f"Session cache saved: {entry_id} ({file_size} bytes, {prompt_tokens} tokens)")
+        except Exception as e:
+            log(f"Failed to save session cache: {e}")
+
+    # -- Eviction --
+
+    def _evict_superseded(self, prompt_text, current_id):
+        """Remove entries whose prompt is a strict prefix of *prompt_text*."""
+        to_remove = [
+            eid for eid, e in self.cache_map["entries"].items()
+            if eid != current_id and prompt_text.startswith(e.get("promptText", ""))
+            and e.get("promptText", "") != prompt_text
+        ]
+        for eid in to_remove:
+            log(f"Session cache evict superseded: {eid}")
+            self._remove_entry(eid)
+
+    def _evict_if_needed(self):
+        entries = self.cache_map.get("entries", {})
+        sorted_ids = sorted(
+            entries.keys(), key=lambda k: entries[k].get("lastAccessedAt", 0),
+        )
+        while (
+            len(entries) > self.max_entries
+            or self.cache_map.get("totalSize", 0) > self.max_size_bytes
+        ):
+            if not sorted_ids:
+                break
+            oldest_id = sorted_ids.pop(0)
+            log(f"Session cache evict LRU: {oldest_id}")
+            self._remove_entry(oldest_id)
+
+    def _remove_entry(self, entry_id):
+        entry = self.cache_map["entries"].get(entry_id)
+        if not entry:
+            return
+        fp = entry.get("filePath")
+        if fp:
+            try:
+                os.unlink(fp)
+            except OSError:
+                pass
+        self.cache_map["totalSize"] -= entry.get("fileSize", 0)
+        del self.cache_map["entries"][entry_id]
+        if self._mem_cache_key == entry_id:
+            self._mem_cache = None
+            self._mem_cache_key = None
+
+    def clear_memory(self):
+        """Drop the in-memory cache (e.g. on model release)."""
+        self._mem_cache = None
+        self._mem_cache_key = None
+
+
+def send(msg):
+    """Send a JSON message to stdout (thread-safe)."""
+    with _output_lock:
+        sys.stdout.write(json.dumps(msg, ensure_ascii=False) + "\n")
+        sys.stdout.flush()
+
+
+def log(msg):
+    """Log to stderr."""
+    sys.stderr.write(f"[mlx-bridge] {msg}\n")
+    sys.stderr.flush()
+
+
+def is_cancelled(req_id):
+    with _cancel_lock:
+        return req_id in _cancel_flags
+
+
+def clear_cancel(req_id):
+    with _cancel_lock:
+        _cancel_flags.discard(req_id)
+
+
+def _get_tokenizer():
+    """Return the underlying tokenizer from the processor."""
+    if _processor is None:
+        return None
+    return getattr(_processor, "tokenizer", _processor)
+
+
+def handle_health(req_id, params):
+    versions = {}
+    try:
+        import mlx_vlm
+        versions["mlx_vlm"] = mlx_vlm.__version__
+    except ImportError:
+        pass
+    try:
+        import mlx_lm
+        versions["mlx_lm"] = mlx_lm.__version__
+    except ImportError:
+        pass
+    if versions:
+        send({"id": req_id, "result": {"ok": True, **versions}})
+    else:
+        send({"id": req_id, "error": {"message": "Neither mlx_vlm nor mlx_lm available"}})
+
+
+def handle_load(req_id, params):
+    global _model, _processor, _config, _model_id, _is_vlm
+
+    model_id = params.get("model")
+    if not model_id:
+        send({"id": req_id, "error": {"message": "Missing 'model' parameter"}})
+        return
+
+    # Release previous model if any
+    _model = None
+    _processor = None
+    _config = None
+    _model_id = None
+    _is_vlm = False
+
+    # vlm param: "auto" (default) tries vlm then lm, True forces vlm, False forces lm
+    vlm_mode = params.get("vlm", "auto")
+
+    log(f"Loading model: {model_id} (vlm={vlm_mode})")
+
+    load_kwargs = {}
+    if params.get("adapter_path"):
+        load_kwargs["adapter_path"] = params["adapter_path"]
+    if params.get("revision"):
+        load_kwargs["revision"] = params["revision"]
+
+    def _load_via_vlm():
+        global _model, _processor, _is_vlm
+        from mlx_vlm import load as vlm_load
+        _model, _processor = vlm_load(model_id, **load_kwargs)
+        _is_vlm = True
+        log(f"Model loaded via mlx-vlm: {model_id}")
+
+    def _load_via_lm():
+        global _model, _processor, _is_vlm
+        from mlx_lm import load as lm_load
+
+        lm_kwargs = dict(load_kwargs)
+        tokenizer_config = params.get("tokenizer_config") or {}
+        model_config = params.get("model_config") or {}
+        if tokenizer_config:
+            lm_kwargs["tokenizer_config"] = tokenizer_config
+        if model_config:
+            lm_kwargs["model_config"] = model_config
+
+        model, tokenizer = lm_load(model_id, **lm_kwargs)
+        _model = model
+        _processor = tokenizer
+        _is_vlm = False
+        log(f"Model loaded via mlx-lm: {model_id}")
+
+    if vlm_mode is True or vlm_mode == "true":
+        _load_via_vlm()
+    elif vlm_mode is False or vlm_mode == "false":
+        _load_via_lm()
+    else:
+        # auto: try vlm first, fall back to lm
+        try:
+            _load_via_vlm()
+        except Exception as vlm_err:
+            log(f"mlx-vlm load failed ({vlm_err}), falling back to mlx-lm")
+            _load_via_lm()
+
+    _config = getattr(_model, "config", {})
+    _model_id = model_id
+
+    send({"id": req_id, "result": {"loaded": True, "model": model_id, "vlm": _is_vlm}})
+
+
+def handle_generate(req_id, params):
+    if _model is None or _processor is None:
+        send({"id": req_id, "error": {"message": "No model loaded"}})
+        return
+
+    from mlx_lm.cache_prompt import make_prompt_cache
+
+    prompt = params.get("prompt", "")
+    max_tokens = params.get("max_tokens", 256)
+    images = params.get("image") or None
+
+    # Sampling parameters
+    kwargs = {}
+    if _is_vlm:
+        # mlx-vlm uses "temperature" directly
+        if "temperature" in params:
+            kwargs["temperature"] = params["temperature"]
+        if "top_p" in params:
+            kwargs["top_p"] = params["top_p"]
+        if "top_k" in params:
+            kwargs["top_k"] = params["top_k"]
+        if "min_p" in params:
+            kwargs["min_p"] = params["min_p"]
+        if "repetition_penalty" in params:
+            kwargs["repetition_penalty"] = params["repetition_penalty"]
+        if "repetition_context_size" in params:
+            kwargs["repetition_context_size"] = params["repetition_context_size"]
+    else:
+        # mlx-lm uses sampler + logits_processors (not raw kwargs)
+        from mlx_lm.sample_utils import make_sampler, make_logits_processors
+        sampler_kwargs = {}
+        if "temperature" in params:
+            sampler_kwargs["temp"] = params["temperature"]
+        if "top_p" in params:
+            sampler_kwargs["top_p"] = params["top_p"]
+        if "top_k" in params:
+            sampler_kwargs["top_k"] = params["top_k"]
+        if "min_p" in params:
+            sampler_kwargs["min_p"] = params["min_p"]
+        if sampler_kwargs:
+            kwargs["sampler"] = make_sampler(**sampler_kwargs)
+        lp_kwargs = {}
+        if "repetition_penalty" in params:
+            lp_kwargs["repetition_penalty"] = params["repetition_penalty"]
+        if "repetition_context_size" in params:
+            lp_kwargs["repetition_context_size"] = params["repetition_context_size"]
+        if lp_kwargs:
+            kwargs["logits_processors"] = make_logits_processors(**lp_kwargs)
+    if "seed" in params:
+        kwargs["seed"] = params["seed"]
+    if "max_kv_size" in params:
+        kwargs["max_kv_size"] = params["max_kv_size"]
+
+    stop_sequences = params.get("stop") or []
+    if isinstance(stop_sequences, str):
+        stop_sequences = [stop_sequences]
+
+    # Session cache: try to reuse a cached prompt prefix (text-only)
+    prompt_cache = None
+    if _session_cache and _session_cache.cache_map is not None and not images:
+        match = _session_cache.find_best_prefix(prompt)
+        if match:
+            loaded = _session_cache.load_cache_for(match)
+            if loaded:
+                prompt_cache = loaded
+
+    if prompt_cache is None:
+        prompt_cache = make_prompt_cache(_model)
+
+    full_text = ""
+    last_response = None
+
+    if _is_vlm:
+        from mlx_vlm import stream_generate
+        gen_iter = stream_generate(
+            model=_model,
+            processor=_processor,
+            prompt=prompt,
+            image=images,
+            max_tokens=max_tokens,
+            prompt_cache=prompt_cache,
+            **kwargs,
+        )
+    else:
+        from mlx_lm import stream_generate
+        gen_iter = stream_generate(
+            _model,
+            _processor,
+            prompt=prompt,
+            max_tokens=max_tokens,
+            prompt_cache=prompt_cache,
+            **kwargs,
+        )
+
+    for response in gen_iter:
+        if is_cancelled(req_id):
+            clear_cancel(req_id)
+            send({
+                "id": req_id,
+                "event": "result",
+                "data": {
+                    "text": full_text,
+                    "interrupted": True,
+                    "prompt_tokens": getattr(response, "prompt_tokens", 0),
+                    "prompt_tps": getattr(response, "prompt_tps", 0),
+                    "generation_tokens": getattr(response, "generation_tokens", 0),
+                    "generation_tps": getattr(response, "generation_tps", 0),
+                    "peak_memory": getattr(response, "peak_memory", 0),
+                },
+            })
+            return
+
+        full_text += response.text
+        last_response = response
+
+        send({
+            "id": req_id,
+            "event": "token",
+            "data": {
+                "token": response.text,
+                "token_id": getattr(response, "token", None),
+            },
+        })
+
+        # Check stop sequences
+        should_stop = False
+        for seq in stop_sequences:
+            if seq and seq in full_text:
+                should_stop = True
+                # Trim text after stop sequence
+                idx = full_text.index(seq)
+                full_text = full_text[:idx]
+                break
+        if should_stop:
+            break
+
+    # Send final result
+    result_data = {
+        "text": full_text,
+        "interrupted": False,
+    }
+    if last_response:
+        result_data.update({
+            "prompt_tokens": getattr(last_response, "prompt_tokens", 0),
+            "prompt_tps": getattr(last_response, "prompt_tps", 0),
+            "generation_tokens": getattr(last_response, "generation_tokens", 0),
+            "generation_tps": getattr(last_response, "generation_tps", 0),
+            "peak_memory": getattr(last_response, "peak_memory", 0),
+            "finish_reason": getattr(last_response, "finish_reason", None),
+        })
+
+    send({"id": req_id, "event": "result", "data": result_data})
+
+    # Session cache: save prompt-only KV state for future reuse (text-only)
+    if _session_cache and last_response and not images and not result_data.get("interrupted"):
+        try:
+            prompt_tokens = getattr(last_response, "prompt_tokens", 0) or 0
+            gen_tokens = getattr(last_response, "generation_tokens", 0) or 0
+            _session_cache.save_after_generation(
+                prompt, prompt_tokens, prompt_cache, gen_tokens,
+            )
+        except Exception as e:
+            log(f"Session cache save error: {e}")
+
+
+def handle_tokenize(req_id, params):
+    tokenizer = _get_tokenizer()
+    if tokenizer is None:
+        send({"id": req_id, "error": {"message": "No model loaded"}})
+        return
+
+    text = params.get("text", "")
+    tokens = tokenizer.encode(text)
+    # Ensure tokens are plain Python ints
+    send({"id": req_id, "result": {"tokens": [int(t) for t in tokens]}})
+
+
+def handle_detokenize(req_id, params):
+    tokenizer = _get_tokenizer()
+    if tokenizer is None:
+        send({"id": req_id, "error": {"message": "No model loaded"}})
+        return
+
+    tokens = params.get("tokens", [])
+    text = tokenizer.decode(tokens)
+    send({"id": req_id, "result": {"text": text}})
+
+
+def handle_apply_chat_template(req_id, params):
+    if _processor is None:
+        send({"id": req_id, "error": {"message": "No model loaded"}})
+        return
+
+    messages = params.get("messages", [])
+    add_generation_prompt = params.get("add_generation_prompt", True)
+
+    # Reserved keys that are not chat template kwargs
+    reserved_keys = {"messages", "add_generation_prompt"}
+    template_kwargs = {k: v for k, v in params.items() if k not in reserved_keys}
+
+    if _is_vlm:
+        # VLM: prefer processor.apply_chat_template (handles image tokens)
+        apply_fn = getattr(_processor, "apply_chat_template", None)
+        if apply_fn is None:
+            apply_fn = _get_tokenizer().apply_chat_template
+    else:
+        # Text-only: use tokenizer directly
+        apply_fn = _get_tokenizer().apply_chat_template
+
+    text = apply_fn(
+        messages,
+        tokenize=False,
+        add_generation_prompt=add_generation_prompt,
+        **template_kwargs,
+    )
+    send({"id": req_id, "result": {"text": text}})
+
+
+def handle_configure_cache(req_id, params):
+    global _session_cache
+
+    enabled = params.get("enabled", True)
+    if not enabled:
+        _session_cache = None
+        send({"id": req_id, "result": {"configured": True, "enabled": False}})
+        return
+
+    cache_dir = os.path.abspath(params.get("cache_dir", os.path.join(
+        os.path.expanduser("~"), ".buttress", "mlx-session-cache",
+    )))
+    max_entries = params.get("max_entries", 100)
+    max_size_bytes = params.get("max_size_bytes", 5 * 1024 * 1024 * 1024)
+
+    _session_cache = SessionCacheManager(cache_dir, max_entries, max_size_bytes)
+    _session_cache.initialize()
+    send({"id": req_id, "result": {
+        "configured": True,
+        "enabled": True,
+        "entries": len(_session_cache.cache_map.get("entries", {})),
+    }})
+
+
+def handle_release(req_id, params):
+    global _model, _processor, _config, _model_id, _is_vlm
+    _model = None
+    _processor = None
+    _config = None
+    _model_id = None
+    _is_vlm = False
+    if _session_cache:
+        _session_cache.clear_memory()
+    log("Model released")
+    send({"id": req_id, "result": {"released": True}})
+
+
+def handle_get_info(req_id, params):
+    info = {"model": _model_id, "loaded": _model is not None}
+    if _model is not None:
+        try:
+            import mlx.core as mx
+            info["peak_memory"] = mx.metal.get_peak_memory() / (1024 * 1024)  # MB
+            info["active_memory"] = mx.metal.get_active_memory() / (1024 * 1024)  # MB
+        except Exception:
+            pass
+    send({"id": req_id, "result": info})
+
+
+# Method dispatch table
+SYNC_METHODS = {
+    "health": handle_health,
+    "load": handle_load,
+    "tokenize": handle_tokenize,
+    "detokenize": handle_detokenize,
+    "apply_chat_template": handle_apply_chat_template,
+    "configure_cache": handle_configure_cache,
+    "release": handle_release,
+    "get_info": handle_get_info,
+}
+
+
+def handle_request(req):
+    req_id = req.get("id")
+    method = req.get("method")
+    params = req.get("params", {})
+
+    if method == "cancel":
+        target_id = params.get("request_id")
+        if target_id:
+            with _cancel_lock:
+                _cancel_flags.add(target_id)
+        return
+
+    handler = SYNC_METHODS.get(method)
+    if handler:
+        try:
+            handler(req_id, params)
+        except Exception as e:
+            log(f"Error in {method}: {traceback.format_exc()}")
+            send({"id": req_id, "error": {"message": str(e)}})
+        return
+
+    if method == "generate":
+        # Run generation in a thread so stdin remains readable for cancel
+        thread = threading.Thread(
+            target=_run_generate_safe, args=(req_id, params), daemon=True,
+        )
+        thread.start()
+        return
+
+    send({"id": req_id, "error": {"message": f"Unknown method: {method}"}})
+
+
+def _run_generate_safe(req_id, params):
+    try:
+        handle_generate(req_id, params)
+    except Exception as e:
+        log(f"Error in generate: {traceback.format_exc()}")
+        send({"id": req_id, "error": {"message": str(e)}})
+    finally:
+        clear_cancel(req_id)
+
+
+def main():
+    log("Bridge started")
+    send({"id": "__init__", "result": {"ready": True, "pid": os.getpid()}})
+
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            req = json.loads(line)
+        except json.JSONDecodeError as e:
+            send({"error": {"message": f"Invalid JSON: {e}"}})
+            continue
+
+        handle_request(req)
+
+    log("Bridge stdin closed, exiting")
+
+
+if __name__ == "__main__":
+    main()