code/models/autoports/google_embeddinggemma_2/server/engine.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import base64
import binascii
import collections
import io
import os
import sys
import threading
import time
import urllib.request
from concurrent.futures import Future
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
from PIL import Image

HF_MODEL = "google/embeddinggemma-2"
HF_REVISION = "914f7f89142e33e77833254d9c9b90c3cef7303b"
SERVED_MODEL = "google/embeddinggemma-2"
MODEL_ALIASES = (SERVED_MODEL, "embeddinggemma-2", "tt-hous/embeddinggemma-2-p150")
DIMENSIONS = (768, 512, 256, 128)
INPUT_TYPES = ("query", "document")
SAMPLING_RATE = 16000
MAX_INPUTS = 2048
MAX_MEDIA_BYTES = 32 * 1024 * 1024
MAX_VIDEO_FRAMES = 1024
URL_TIMEOUT_S = 15
L1_SMALL_SIZE = 32768
TRACE_REGION_SIZE = 2_500_000_000
MEDIA_KEYS = ("image", "audio", "video")


class RequestError(ValueError):
    def __init__(self, message, param=None, status=400, code=None):
        super().__init__(message)
        self.message = message
        self.param = param
        self.status = status
        self.code = code


def parse_mesh_shape(text):
    parts = [p for p in text.lower().replace(",", "x").split("x") if p.strip()]
    if len(parts) != 2 or not all(p.strip().isdigit() for p in parts):
        raise ValueError(f"EMBEDDINGGEMMA_MESH_SHAPE must look like 1x1, got {text!r}")
    return int(parts[0]), int(parts[1])


def resolve_snapshot(model, revision):
    if Path(model).is_dir():
        return Path(model)
    from huggingface_hub import snapshot_download

    try:
        return Path(snapshot_download(model, revision=revision, local_files_only=True))
    except Exception:
        return Path(snapshot_download(model, revision=revision))


def strip_data_url(value):
    if value.startswith("data:"):
        head, sep, data = value.partition(",")
        if not sep or ";base64" not in head:
            raise RequestError("data URLs must be base64 encoded")
        return data
    return value


def fetch_url(url):
    req = urllib.request.Request(url, headers={"User-Agent": "embeddinggemma-2-tt-server"})
    try:
        with urllib.request.urlopen(req, timeout=URL_TIMEOUT_S) as resp:
            data = resp.read(MAX_MEDIA_BYTES + 1)
    except Exception as exc:
        raise RequestError(f"could not fetch {url}: {exc}")
    if len(data) > MAX_MEDIA_BYTES:
        raise RequestError(f"{url} is larger than {MAX_MEDIA_BYTES} bytes")
    return data


def media_bytes(value, what):
    if isinstance(value, dict) and isinstance(value.get("url"), str):
        value = value["url"]
    if not isinstance(value, str) or not value:
        raise RequestError(f"{what} must be a non-empty base64 string or URL")
    if value.startswith(("http://", "https://")):
        if os.environ.get("EMBEDDINGGEMMA_ALLOW_URLS", "1") == "0":
            raise RequestError(f"{what}: URL inputs are disabled on this server; send base64")
        return fetch_url(value)
    try:
        data = base64.b64decode(strip_data_url(value), validate=True)
    except (binascii.Error, ValueError):
        raise RequestError(f"{what} is not valid base64 and not an http(s) URL")
    if not data:
        raise RequestError(f"{what} is empty")
    if len(data) > MAX_MEDIA_BYTES:
        raise RequestError(f"{what} is larger than {MAX_MEDIA_BYTES} bytes")
    return data


def decode_image(value, what="image"):
    data = media_bytes(value, what)
    try:
        img = Image.open(io.BytesIO(data))
        img.load()
    except Exception as exc:
        raise RequestError(f"{what} is not a decodable image: {exc}")
    return img


def decode_audio(value, what="audio"):
    import soundfile as sf

    data = media_bytes(value, what)
    try:
        wav, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False)
    except Exception as exc:
        raise RequestError(f"{what} is not a decodable WAV file: {exc}")
    if sr != SAMPLING_RATE:
        raise RequestError(f"{what} sample rate is {sr} Hz; it must be {SAMPLING_RATE} Hz")
    if wav.ndim == 2:
        wav = wav.mean(axis=1, dtype=np.float32)
    if wav.size == 0:
        raise RequestError(f"{what} has no samples")
    return {"array": np.ascontiguousarray(wav, dtype=np.float32), "sampling_rate": sr}


def decode_video(value, what="video"):
    if isinstance(value, list):
        value = {"frames": value}
    if not isinstance(value, dict) or not isinstance(value.get("frames"), list) or not value["frames"]:
        raise RequestError(f"{what} must be an object with a non-empty 'frames' list")
    if len(value["frames"]) > MAX_VIDEO_FRAMES:
        raise RequestError(f"{what} has {len(value['frames'])} frames; the limit is {MAX_VIDEO_FRAMES}")
    frames = [np.asarray(decode_image(f, f"{what}.frames[{i}]").convert("RGB")) for i, f in enumerate(value["frames"])]
    if len({f.shape for f in frames}) != 1:
        raise RequestError(f"{what} frames must all have the same size")
    arr = np.stack(frames)
    fps = value.get("fps")
    metadata = None
    if fps is not None:
        if not isinstance(fps, (int, float)) or isinstance(fps, bool) or fps <= 0:
            raise RequestError(f"{what}.fps must be a positive number")
        metadata = {"fps": float(fps), "total_num_frames": len(frames), "duration": len(frames) / float(fps)}
    return {"array": arr, "video_metadata": metadata}


def decode_input(item, index):
    what = f"input[{index}]"
    if isinstance(item, str):
        if not item:
            raise RequestError(f"{what} is an empty string", param="input")
        return item
    if not isinstance(item, dict):
        raise RequestError(f"{what} must be a string or an object with text/image/audio/video fields", param="input")
    bad = sorted(set(item) - {"text", *MEDIA_KEYS})
    if bad:
        raise RequestError(f"{what} has unsupported fields {bad}", param="input")
    media = [k for k in MEDIA_KEYS if item.get(k) is not None]
    text = item.get("text")
    if text is not None and not isinstance(text, str):
        raise RequestError(f"{what}.text must be a string", param="input")
    if len(media) > 1:
        raise RequestError(f"{what} mixes {media}; one media kind per input is supported", param="input")
    if not media:
        if not text:
            raise RequestError(f"{what} has no text, image, audio or video", param="input")
        return text
    kind = media[0]
    decoder = {"image": decode_image, "audio": decode_audio, "video": decode_video}[kind]
    out = {}
    for key in item:
        if key == kind:
            out[kind] = decoder(item[kind], f"{what}.{kind}")
        elif key == "text" and text:
            out["text"] = text
    return out


def normalize_inputs(raw):
    if isinstance(raw, str):
        raw = [raw]
    if not isinstance(raw, list):
        raise RequestError("input must be a string or a list of strings or objects", param="input")
    if not raw:
        raise RequestError("input must not be empty", param="input")
    if len(raw) > MAX_INPUTS:
        raise RequestError(f"input has {len(raw)} items; the limit is {MAX_INPUTS}", param="input")
    if all(isinstance(x, int) and not isinstance(x, bool) for x in raw) or any(isinstance(x, list) for x in raw):
        raise RequestError("token id inputs are not supported; send strings", param="input")
    return [decode_input(x, i) for i, x in enumerate(raw)]


@dataclass
class Job:
    items: list
    dim: int
    key: tuple = ()
    future: Future = field(default_factory=Future)
    enqueued: float = field(default_factory=time.perf_counter)


class Engine:
    def __init__(self, embedder, mesh_device, policy_id=None, max_batch_items=32, batch_wait_ms=0.0):
        self.embedder = embedder
        self.mesh_device = mesh_device
        self.policy_id = policy_id
        self.max_batch_items = max_batch_items
        self.batch_wait_s = batch_wait_ms / 1000.0
        self.device_lock = threading.Lock()
        self.text_lock = threading.Lock()
        self.media_lock = threading.Lock()
        self.bucket_grouping = os.environ.get("EMBEDDINGGEMMA_BUCKET_GROUPING", "1") != "0"
        self.pending = collections.deque()
        self.cv = threading.Condition()
        self.closing = False
        sizes = {b for b, _, _ in embedder.traces.slots} if embedder.traces is not None else set()
        self.slot_sizes = sorted(b for b in sizes if b <= max_batch_items) or [max_batch_items]
        self.slot_cost = {}
        self.record = None
        self.stats = {
            "requests": 0,
            "inputs": 0,
            "device_batches": 0,
            "encode_groups": 0,
            "max_jobs_per_batch": 0,
            "device_busy_s": 0.0,
            "device_thread_errors": 0,
        }
        self.thread = threading.Thread(target=self._loop, name="embeddinggemma-device", daemon=True)
        self.started = time.time()
        self.thread.start()

    @property
    def max_tokens(self):
        return self.embedder.max_tokens

    def prompt_for(self, input_type):
        if input_type is None:
            return None
        if input_type not in self.embedder.prompts:
            known = sorted(set(INPUT_TYPES) | set(self.embedder.prompts))
            raise RequestError(f"unknown input_type {input_type!r}; known: {known}", param="input_type")
        return self.embedder.prompts[input_type]

    def prepare(self, inputs, input_type=None):
        from models.autoports.google_embeddinggemma_2.tt.multimodal import item_kind

        prompt = self.prompt_for(input_type)
        emb = self.embedder
        out = []
        for i, x in enumerate(inputs):
            try:
                if isinstance(x, str):
                    with self.text_lock:
                        ids = emb.tokenizer([(prompt or "") + x], add_special_tokens=True)["input_ids"][0]
                    it = {"input_ids": ids}
                else:
                    with self.media_lock:
                        it = emb.process([x], prompt=prompt or "")[0]
            except RequestError:
                raise
            except ValueError as exc:
                raise RequestError(f"input[{i}]: {exc}", param="input")
            n = len(it["input_ids"])
            if emb.max_tokens is not None and n > emb.max_tokens:
                raise RequestError(
                    f"input[{i}] has {n} tokens; the maximum is {emb.max_tokens}",
                    param="input",
                    code="context_length_exceeded",
                )
            kind = item_kind(it, emb.pipe.tokens)
            if kind is not None and emb.traces is not None and emb.media_slot_for(kind, [it]) is None:
                raise RequestError(
                    f"input[{i}] ({kind}, {n} tokens) does not fit any captured {kind} trace slot",
                    param="input",
                )
            out.append(it)
        return out

    def item_key(self, it):
        from models.autoports.google_embeddinggemma_2.tt.model import uniform_physical_length
        from models.autoports.google_embeddinggemma_2.tt.multimodal import item_kind, item_patches

        kind = item_kind(it, self.embedder.pipe.tokens)
        if kind is not None:
            return ("media", kind, item_patches(it))
        return ("text", uniform_physical_length(self.embedder.seq_len_for([len(it["input_ids"])])))

    def submit(self, items, dim):
        if dim not in DIMENSIONS:
            raise RequestError(f"dimensions must be one of {list(DIMENSIONS)}", param="dimensions")
        job = Job(items, dim, tuple(sorted({self.item_key(it) for it in items})))
        with self.cv:
            self.pending.append(job)
            self.cv.notify()
        return job.future

    def queue_depth(self):
        return len(self.pending)

    def run_direct(self, fn):
        with self.device_lock:
            return fn(self.embedder)

    def target_items(self, n, phys=None):
        if n >= self.slot_sizes[-1]:
            return max(n, self.slot_sizes[-1]) if n > self.max_batch_items else self.slot_sizes[-1]
        up = min(b for b in self.slot_sizes if b >= n)
        down = max((b for b in self.slot_sizes if b <= n), default=n)
        cost_up, cost_down = self.slot_cost.get((up, phys)), self.slot_cost.get((down, phys))
        if cost_up is not None and cost_down is not None:
            return n if cost_up / n <= cost_down / down else down
        return n if up - n <= n - down else down

    def text_phys(self, jobs):
        from models.autoports.google_embeddinggemma_2.tt.model import uniform_physical_length

        items = [it for j in jobs for it in j.items]
        if any("pixel_values" in it or "input_features" in it for it in items):
            return None
        return uniform_physical_length(self.embedder.seq_len_for([max(len(it["input_ids"]) for it in items)]))

    def _take(self):
        with self.cv:
            while not self.pending and not self.closing:
                self.cv.wait()
            if not self.pending:
                return None
            deadline = time.perf_counter() + self.batch_wait_s
            while not self.closing and sum(len(j.items) for j in self.pending) < self.max_batch_items:
                left = deadline - time.perf_counter()
                if left <= 0:
                    break
                self.cv.wait(left)
            head = self.pending[0].key
            avail, n = [], 0
            for j in self.pending:
                if self.bucket_grouping and j.key != head:
                    continue
                if avail and n + len(j.items) > self.max_batch_items:
                    break
                avail.append(j)
                n += len(j.items)
            phys = head[0][1] if len(head) == 1 and head[0][0] == "text" else None
            k = self.target_items(n, phys if self.bucket_grouping else self.text_phys(avail))
            jobs, c = [], 0
            for j in avail:
                if jobs and c + len(j.items) > k:
                    break
                jobs.append(j)
                c += len(j.items)
            taken = {id(j) for j in jobs}
            self.pending = collections.deque(j for j in self.pending if id(j) not in taken)
            return jobs

    def _encode(self, jobs):
        dims = tuple(sorted({j.dim for j in jobs}, reverse=True))
        rows = [(j, k, it) for j in jobs for k, it in enumerate(j.items)]
        groups = {}
        for r in rows:
            key = self.item_key(r[2]) if self.bucket_grouping else ()
            groups.setdefault(key, []).append(r)
        results = {id(j): np.zeros((len(j.items), j.dim), dtype=np.float32) for j in jobs}
        for key, grp in groups.items():
            items = [it for _, _, it in grp]
            with self.device_lock:
                t0 = time.perf_counter()
                out = self.embedder.encode_processed(items, dims=dims)
                self.stats["device_busy_s"] += time.perf_counter() - t0
            self.stats["encode_groups"] += 1
            if self.record is not None:
                self.record.append(
                    {
                        "key": key,
                        "items": items,
                        "dims": dims,
                        "row_dims": [j.dim for j, _, _ in grp],
                        "out": out,
                        "n_jobs": len(jobs),
                    }
                )
            for pos, (j, k, _) in enumerate(grp):
                results[id(j)][k] = out[j.dim][pos]
        for j in jobs:
            self._resolve(j, results[id(j)])

    @staticmethod
    def _resolve(job, value=None, exc=None):
        if job.future.done():
            return
        if exc is not None:
            job.future.set_exception(exc)
        else:
            job.future.set_result(value)

    def _loop(self):
        while True:
            jobs = self._take()
            if jobs is None:
                return
            try:
                self.stats["device_batches"] += 1
                self.stats["max_jobs_per_batch"] = max(self.stats["max_jobs_per_batch"], len(jobs))
                try:
                    self._encode(jobs)
                except Exception as exc:
                    if len(jobs) == 1:
                        self._resolve(jobs[0], exc=exc)
                        continue
                    for j in jobs:
                        try:
                            self._encode([j])
                        except Exception as one:
                            self._resolve(j, exc=one)
            except Exception as exc:
                self.stats["device_thread_errors"] += 1
                for j in jobs:
                    try:
                        self._resolve(j, exc=exc)
                    except Exception:
                        pass

    def alive(self):
        return self.thread.is_alive()

    def trace_counters(self):
        traces = self.embedder.traces
        if traces is None:
            return {"traced": False}
        slots = {f"text:{b}x{p}": s for (b, p, _), s in traces.slots.items()}
        slots.update(
            {f"{s.shape.kind}:{s.shape.batch}x{s.shape.physical_len}:tower{s.shape.tower_n}": s for s in traces.extra}
        )
        per = {
            k: {"trace_id": str(s.trace_id), "replays": s.replays, "refreshes": s.refreshes} for k, s in slots.items()
        }
        return {
            "captures": traces.captures,
            "live_traces": sum(1 for s in slots.values() if s.trace_id is not None),
            "replays_total": sum(s.replays for s in slots.values()),
            "slots": per,
        }

    def warm(self):
        traces = self.embedder.traces
        if traces is None:
            return 0
        import ttnn

        slots = list(traces.slots.values()) + list(traces.extra)
        with self.device_lock:
            for s in slots:
                s.replay()
            ttnn.synchronize_device(self.mesh_device)
            for (batch, phys, exact), s in traces.slots.items():
                if exact:
                    continue
                t0 = time.perf_counter()
                s.replay()
                ttnn.synchronize_device(self.mesh_device)
                self.slot_cost[(batch, phys)] = time.perf_counter() - t0
        return len(slots)

    def slot_summary(self):
        traces = self.embedder.traces
        if traces is None:
            return {"traced": False}
        text = sorted(traces.slots)
        media = [s.shape for s in traces.extra]
        return {
            "traced": True,
            "text_slots": len(text),
            "media_slots": len(media),
            "text_buckets": sorted({p for _, p, _ in text}),
            "text_batches": sorted({b for b, _, _ in text}),
            "media": [f"{m.kind}:b{m.batch}:len{m.physical_len}:tower{m.tower_n}" for m in media],
        }

    def close(self):
        with self.cv:
            self.closing = True
            self.cv.notify_all()
        self.thread.join(timeout=60)


def build_engine():
    import ttnn
    from models.autoports.google_embeddinggemma_2.tt import precision_policy as PP
    from models.autoports.google_embeddinggemma_2.tt.embedder import build_multimodal_embedder

    sys.setswitchinterval(float(os.environ.get("EMBEDDINGGEMMA_SWITCH_INTERVAL_MS", "0.5")) / 1000.0)
    rows, cols = parse_mesh_shape(os.environ.get("EMBEDDINGGEMMA_MESH_SHAPE", "1x1"))
    if (rows, cols) != (1, 1):
        raise ValueError(f"the embedder runs on one p150 (mesh 1x1); EMBEDDINGGEMMA_MESH_SHAPE={rows}x{cols}")
    snapshot = resolve_snapshot(os.environ.get("HF_MODEL", HF_MODEL), os.environ.get("HF_REVISION", HF_REVISION))
    mesh = ttnn.open_mesh_device(
        ttnn.MeshShape(rows, cols), l1_small_size=L1_SMALL_SIZE, trace_region_size=TRACE_REGION_SIZE
    )
    try:
        policy = PP.load()
        embedder = build_multimodal_embedder(snapshot, mesh, traced=True, precision_policy=policy)
        report = PP.runtime_report(embedder)
        if not report["consumed"]:
            raise RuntimeError(f"precision policy {policy.id} not applied: {report['mismatches'][:4]}")
    except Exception:
        ttnn.close_mesh_device(mesh)
        raise
    g = mesh.compute_with_storage_grid_size()
    engine = Engine(
        embedder,
        mesh,
        policy_id=policy.id,
        max_batch_items=int(os.environ.get("EMBEDDINGGEMMA_MAX_BATCH_ITEMS", "32")),
        batch_wait_ms=float(os.environ.get("EMBEDDINGGEMMA_BATCH_WAIT_MS", "0")),
    )
    engine.snapshot = str(snapshot)
    engine.grid = f"{g.x}x{g.y}"
    engine.mesh_shape = f"{rows}x{cols}"
    return engine


def shutdown_engine(engine):
    import ttnn

    engine.close()
    with engine.device_lock:
        engine.embedder.release()
        ttnn.close_mesh_device(engine.mesh_device)