| |
|
|
| |
|
|
| import asyncio |
| import base64 |
| import logging |
| import os |
| import sys |
| import time |
| from contextlib import asynccontextmanager |
| from typing import Any, Optional |
|
|
| import numpy as np |
| import orjson |
| from fastapi import FastAPI, Request |
| from fastapi.concurrency import run_in_threadpool |
| from fastapi.exceptions import RequestValidationError |
| from fastapi.responses import JSONResponse, Response |
| from pydantic import BaseModel, StrictInt |
| from starlette.exceptions import HTTPException as StarletteHTTPException |
|
|
| from models.autoports.google_embeddinggemma_2.server.engine import ( |
| DIMENSIONS, |
| MODEL_ALIASES, |
| SERVED_MODEL, |
| RequestError, |
| build_engine, |
| normalize_inputs, |
| shutdown_engine, |
| ) |
|
|
| log = logging.getLogger("embeddinggemma2.server") |
| INLINE_TEXTS = 8 |
|
|
|
|
| class EmbeddingRequest(BaseModel): |
| input: Any = None |
| model: Optional[str] = None |
| dimensions: Optional[StrictInt] = None |
| encoding_format: Optional[str] = "float" |
| input_type: Optional[str] = None |
| user: Optional[str] = None |
|
|
|
|
| def error_response(status, message, param=None, code=None, kind="invalid_request_error"): |
| return JSONResponse( |
| status_code=status, content={"error": {"message": message, "type": kind, "param": param, "code": code}} |
| ) |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(app): |
| t0 = time.perf_counter() |
| engine = build_engine() |
| built = time.perf_counter() |
| warmed = engine.warm() |
| engine.startup_s = {"build_and_capture": round(built - t0, 1), "warm": round(time.perf_counter() - built, 1)} |
| if DEMO: |
| from models.autoports.google_embeddinggemma_2.demo.mount import build_demo_index |
|
|
| t1 = time.perf_counter() |
| await run_in_threadpool(build_demo_index, app, engine) |
| engine.startup_s["demo_index"] = round(time.perf_counter() - t1, 1) |
| app.state.engine = engine |
| log.warning( |
| "embeddinggemma-2 ready: policy=%s mesh=%s grid=%s slots=%s warmed=%d startup=%s", |
| engine.policy_id, |
| engine.mesh_shape, |
| engine.grid, |
| engine.slot_summary(), |
| warmed, |
| engine.startup_s, |
| ) |
| try: |
| yield |
| finally: |
| app.state.engine = None |
| shutdown_engine(engine) |
|
|
|
|
| DEMO = os.environ.get("EMBEDDINGGEMMA_DEMO", "0") == "1" |
| app = FastAPI(title="EmbeddingGemma 2 on Tenstorrent p150", lifespan=lifespan) |
|
|
|
|
| @app.exception_handler(RequestError) |
| async def request_error_handler(request: Request, exc: RequestError): |
| return error_response(exc.status, exc.message, exc.param, exc.code) |
|
|
|
|
| @app.exception_handler(RequestValidationError) |
| async def validation_error_handler(request: Request, exc: RequestValidationError): |
| parts = [] |
| for e in exc.errors(): |
| loc = ".".join(str(x) for x in e.get("loc", ()) if x != "body") |
| parts.append(f"{loc}: {e.get('msg')}" if loc else str(e.get("msg"))) |
| return error_response(400, "; ".join(parts) or "invalid request") |
|
|
|
|
| @app.exception_handler(StarletteHTTPException) |
| async def http_error_handler(request: Request, exc: StarletteHTTPException): |
| return error_response( |
| exc.status_code, str(exc.detail), kind="not_found_error" if exc.status_code == 404 else "error" |
| ) |
|
|
|
|
| @app.exception_handler(Exception) |
| async def internal_error_handler(request: Request, exc: Exception): |
| log.exception("unhandled error") |
| return error_response(500, f"internal error: {type(exc).__name__}: {exc}", kind="server_error") |
|
|
|
|
| def engine_or_503(request): |
| engine = getattr(request.app.state, "engine", None) |
| if engine is None: |
| raise RequestError("model is not loaded", status=503, code="not_ready") |
| if not engine.alive(): |
| raise RequestError("device thread is not running", status=503, code="device_thread_dead") |
| return engine |
|
|
|
|
| def health_payload(engine): |
| return { |
| "status": "ok", |
| "model": SERVED_MODEL, |
| "precision_policy": engine.policy_id, |
| "mesh_shape": engine.mesh_shape, |
| "grid": engine.grid, |
| "max_tokens": engine.max_tokens, |
| "switch_interval_ms": round(sys.getswitchinterval() * 1000, 3), |
| "batching": { |
| "max_batch_items": engine.max_batch_items, |
| "batch_wait_ms": engine.batch_wait_s * 1000, |
| "slot_ms": {f"{b}x{p}": round(t * 1e3, 2) for (b, p), t in sorted(engine.slot_cost.items())}, |
| }, |
| "dimensions": list(DIMENSIONS), |
| "slots": engine.slot_summary(), |
| "startup_s": engine.startup_s, |
| "queue_depth": engine.queue_depth(), |
| "stats": dict(engine.stats), |
| "device_thread_alive": engine.alive(), |
| "bucket_grouping": engine.bucket_grouping, |
| "traces": engine.trace_counters(), |
| "uptime_s": round(time.time() - engine.started, 1), |
| } |
|
|
|
|
| @app.get("/health") |
| async def health(request: Request): |
| return health_payload(engine_or_503(request)) |
|
|
|
|
| @app.get("/v1/models") |
| async def models(request: Request): |
| engine = engine_or_503(request) |
| return { |
| "object": "list", |
| "data": [ |
| { |
| "id": SERVED_MODEL, |
| "object": "model", |
| "created": int(engine.started), |
| "owned_by": "google", |
| "root": SERVED_MODEL, |
| "max_model_len": engine.max_tokens, |
| "dimensions": list(DIMENSIONS), |
| "modalities": ["text", "image", "audio", "video"], |
| } |
| ], |
| } |
|
|
|
|
| def encode_vector(vec, fmt): |
| vec = np.ascontiguousarray(vec, dtype="<f4") |
| if fmt == "base64": |
| return base64.b64encode(vec.tobytes()).decode("ascii") |
| return vec |
|
|
|
|
| @app.post("/v1/embeddings") |
| async def embeddings(body: EmbeddingRequest, request: Request): |
| engine = engine_or_503(request) |
| if body.model is not None and body.model not in MODEL_ALIASES: |
| raise RequestError(f"model {body.model!r} is not served; use {SERVED_MODEL!r}", "model", 404, "model_not_found") |
| dim = 768 if body.dimensions is None else body.dimensions |
| if dim not in DIMENSIONS: |
| raise RequestError(f"dimensions must be one of {list(DIMENSIONS)}, got {dim}", "dimensions") |
| fmt = body.encoding_format or "float" |
| if fmt not in ("float", "base64"): |
| raise RequestError(f"encoding_format must be 'float' or 'base64', got {fmt!r}", "encoding_format") |
| if body.input is None: |
| raise RequestError("input is required", "input") |
| engine.prompt_for(body.input_type) |
| raw = body.input |
| if isinstance(raw, str) or ( |
| isinstance(raw, list) and len(raw) <= INLINE_TEXTS and all(isinstance(x, str) for x in raw) |
| ): |
| items = engine.prepare(normalize_inputs(raw), body.input_type) |
| else: |
| inputs = await run_in_threadpool(normalize_inputs, raw) |
| items = await run_in_threadpool(engine.prepare, inputs, body.input_type) |
| vectors = await asyncio.wrap_future(engine.submit(items, dim)) |
| engine.stats["requests"] += 1 |
| engine.stats["inputs"] += len(items) |
| tokens = sum(len(it["input_ids"]) for it in items) |
| payload = { |
| "object": "list", |
| "data": [ |
| {"object": "embedding", "index": i, "embedding": encode_vector(v, fmt)} for i, v in enumerate(vectors) |
| ], |
| "model": SERVED_MODEL, |
| "usage": {"prompt_tokens": tokens, "total_tokens": tokens}, |
| } |
| return Response(orjson.dumps(payload, option=orjson.OPT_SERIALIZE_NUMPY), media_type="application/json") |
|
|
|
|
| if DEMO: |
| from models.autoports.google_embeddinggemma_2.demo.mount import mount_demo |
|
|
| mount_demo(app, lambda a: health_payload(a.state.engine) if getattr(a.state, "engine", None) else {}) |
|
|
