| |
|
|
| |
|
|
| 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) |
|
|
