add monitoring

app.py

@@ -1,5 +1,5 @@
 # app.py
-import asyncio
+from time import perf_counter
 from typing import List, Optional
 
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
@@ -8,8 +8,7 @@ import uvicorn
 
 from util import get_runner, SmolVLMRunner
 
-
-app = FastAPI(title="SmolVLM Inference API", version="1.0.0")
+app = FastAPI(title="SmolVLM Inference API", version="1.1.0")
 _runner: Optional[SmolVLMRunner] = None
 
 
@@ -40,56 +39,70 @@ async def generate_from_files(
     temperature: Optional[float] = Form(None),
     top_p: Optional[float] = Form(None),
 ):
-    """
-    Multipart form endpoint:
-      - prompt: str
-      - images: one or more image files (image/*)
-    """
     if not images:
         raise HTTPException(status_code=400, detail="At least one image must be provided.")
 
-    # Read all files into memory (simple & fine for moderate sizes)
+    t_req_start = perf_counter()
+
+    # Read files
+    t_load_start = perf_counter()
     blobs = []
     for f in images:
         if not f.content_type or not f.content_type.startswith("image/"):
             raise HTTPException(status_code=415, detail=f"Unsupported file type: {f.content_type}")
         blobs.append(await f.read())
-
     pil_images = _runner.load_pil_from_bytes(blobs)
-    text = _runner.generate(
+    t_load_end = perf_counter()
+
+    text, inner_metrics = _runner.generate(
         prompt=prompt,
         images=pil_images,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         top_p=top_p,
+        return_stats=True,
     )
-    return {"text": text}
+
+    t_req_end = perf_counter()
+    metrics = {
+        **inner_metrics,
+        "request_ms": {
+            "image_load": round((t_load_end - t_load_start) * 1000.0, 2),
+            "end_to_end": round((t_req_end - t_req_start) * 1000.0, 2),
+        },
+    }
+    return {"text": text, "metrics": metrics}
 
 
 @app.post("/generate_urls")
 async def generate_from_urls(req: URLRequest):
-    """
-    JSON endpoint:
-      {
-        "prompt": "...",
-        "image_urls": ["https://...","https://..."],
-        "max_new_tokens": 300,
-        "temperature": 0.2,
-        "top_p": 0.95
-      }
-    """
+    t_req_start = perf_counter()
+
     if len(req.image_urls) == 0:
         raise HTTPException(status_code=400, detail="At least one image URL is required.")
 
+    t_load_start = perf_counter()
     pil_images = _runner.load_pil_from_urls([str(u) for u in req.image_urls])
-    text = _runner.generate(
+    t_load_end = perf_counter()
+
+    text, inner_metrics = _runner.generate(
         prompt=req.prompt,
         images=pil_images,
         max_new_tokens=req.max_new_tokens,
         temperature=req.temperature,
         top_p=req.top_p,
+        return_stats=True,
     )
-    return {"text": text}
+
+    t_req_end = perf_counter()
+    metrics = {
+        **inner_metrics,
+        "request_ms": {
+            "image_load": round((t_load_end - t_load_start) * 1000.0, 2),
+            "end_to_end": round((t_req_end - t_req_start) * 1000.0, 2),
+        },
+    }
+    return {"text": text, "metrics": metrics}
 
 
 if __name__ == "__main__":

util.py

@@ -1,57 +1,85 @@
-
-# util.py (patched cache handling for HF Spaces)
+# util.py (Spaces-safe + metrics)
 import os
 from pathlib import Path
+from time import perf_counter
+import threading
+from io import BytesIO
+from typing import List, Sequence, Tuple, Dict, Any
+
+import torch
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers.image_utils import load_image as hf_load_image
 
-# Put every cache under /tmp (always writable in Spaces)
+# ---- Writable caches (HF Spaces safe) ----
 CACHE_DIR = os.getenv("HF_CACHE_DIR", "/tmp/hf-cache")
 Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)
-
-# Make sure libraries don't fall back to "~/.cache" -> "/.cache"
 os.environ.setdefault("HF_HOME", CACHE_DIR)
 os.environ.setdefault("TRANSFORMERS_CACHE", CACHE_DIR)
 os.environ.setdefault("HUGGINGFACE_HUB_CACHE", CACHE_DIR)
 os.environ.setdefault("XDG_CACHE_HOME", CACHE_DIR)
 os.environ.setdefault("TORCH_HOME", CACHE_DIR)
 
-import threading
-from io import BytesIO
-from typing import List, Sequence
-import torch
-from PIL import Image
-from transformers import AutoProcessor, AutoModelForVision2Seq
-from transformers.image_utils import load_image as hf_load_image
+
+def _has_flash_attn() -> bool:
+    try:
+        import flash_attn  # noqa: F401
+        return True
+    except Exception:
+        return False
+
+
+def _pick_backend_and_dtype():
+    if not torch.cuda.is_available():
+        return "eager", torch.float32, "cpu"
+
+    major, _ = torch.cuda.get_device_capability()
+    dev = "cuda"
+    bf16_ok = torch.cuda.is_bf16_supported()
+    dtype = torch.bfloat16 if bf16_ok else torch.float16
+    if major >= 8:  # Ampere+
+        attn = "flash_attention_2" if _has_flash_attn() else "sdpa"
+    else:
+        attn = "sdpa"
+    return attn, dtype, dev
 
 
 class SmolVLMRunner:
+    """Portable wrapper with per-call metrics."""
+
     def __init__(self, model_id: str | None = None, device: str | None = None):
         self.model_id = model_id or os.getenv("SMOLVLM_MODEL_ID", "HuggingFaceTB/SmolVLM-Instruct")
-        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
 
-        # Use the writable cache dir explicitly
+        attn_impl, dtype, dev = _pick_backend_and_dtype()
+        attn_impl = os.getenv("SMOLVLM_ATTN", attn_impl)  # optional override
+        self.device = device or dev
+        self.dtype = dtype
+        self.attn_impl = attn_impl
+
+        if self.device == "cuda" and self.attn_impl == "sdpa":
+            try:
+                from torch.backends.cuda import sdp_kernel
+                sdp_kernel(enable_flash=False, enable_mem_efficient=True, enable_math=True)
+            except Exception:
+                pass
+
         self.processor = AutoProcessor.from_pretrained(self.model_id, cache_dir=CACHE_DIR)
+        self.model = AutoModelForVision2Seq.from_pretrained(
+            self.model_id,
+            torch_dtype=self.dtype,
+            _attn_implementation=self.attn_impl,
+            cache_dir=CACHE_DIR,
+        ).to(self.device)
 
-        attn_impl = "flash_attention_2" if self.device == "cuda" else "eager"
         try:
-            self.model = AutoModelForVision2Seq.from_pretrained(
-                self.model_id,
-                torch_dtype=self.dtype,
-                _attn_implementation=attn_impl,
-                cache_dir=CACHE_DIR,
-            ).to(self.device)
+            self.model.config._attn_implementation = self.attn_impl
         except Exception:
-            # Fallback if flash-attn isn't available in the environment
-            self.model = AutoModelForVision2Seq.from_pretrained(
-                self.model_id,
-                torch_dtype=self.dtype,
-                _attn_implementation="eager",
-                cache_dir=CACHE_DIR,
-            ).to(self.device)
+            pass
 
         self.model.eval()
         self._lock = threading.Lock()
 
+    # ---------- Image utils ----------
     @staticmethod
     def _ensure_rgb(img: Image.Image) -> Image.Image:
         return img.convert("RGB") if img.mode != "RGB" else img
@@ -64,30 +92,118 @@ class SmolVLMRunner:
     def load_pil_from_bytes(cls, blobs: Sequence[bytes]) -> List[Image.Image]:
         return [cls._ensure_rgb(Image.open(BytesIO(b))) for b in blobs]
 
-    def generate(self, prompt: str, images: Sequence[Image.Image], max_new_tokens: int = 300,
-                 temperature: float | None = None, top_p: float | None = None) -> str:
+    # ---------- Inference ----------
+    def generate(
+        self,
+        prompt: str,
+        images: Sequence[Image.Image],
+        max_new_tokens: int = 300,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        return_stats: bool = False,
+    ) -> str | Tuple[str, Dict[str, Any]]:
+        """
+        Returns str by default.
+        If return_stats=True, returns (text, metrics_dict).
+        """
+        meta = {
+            "model_id": self.model_id,
+            "device": self.device,
+            "dtype": str(self.dtype).replace("torch.", ""),
+            "attn_backend": self.attn_impl,
+            "image_count": len(images),
+            "max_new_tokens": int(max_new_tokens),
+            "temperature": None if temperature is None else float(temperature),
+            "top_p": None if top_p is None else float(top_p),
+        }
+
+        t0 = perf_counter()
         content = [{"type": "image"} for _ in images] + [{"type": "text", "text": prompt}]
         messages = [{"role": "user", "content": content}]
         chat_prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True)
 
+        # Preprocess (tokenize + vision)
         inputs = self.processor(text=chat_prompt, images=list(images), return_tensors="pt")
         inputs = {k: (v.to(self.device) if hasattr(v, "to") else v) for k, v in inputs.items()}
+        t_pre_end = perf_counter()
 
+        # Inference (generate)
         gen_kwargs = dict(max_new_tokens=max_new_tokens)
         if temperature is not None:
             gen_kwargs["temperature"] = float(temperature)
         if top_p is not None:
             gen_kwargs["top_p"] = float(top_p)
 
-        with self._lock, torch.inference_mode():
-            generated_ids = self.model.generate(**inputs, **gen_kwargs)
+        if self.device == "cuda":
+            torch.cuda.synchronize()
+            torch.cuda.reset_peak_memory_stats()
 
-        text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        with self._lock, torch.inference_mode():
+            t_inf_start = perf_counter()
+            out_ids = self.model.generate(**inputs, **gen_kwargs)
+            if self.device == "cuda":
+                torch.cuda.synchronize()
+            t_inf_end = perf_counter()
+
+        # Decode
+        text = self.processor.batch_decode(out_ids, skip_special_tokens=True)[0].strip()
         if text.startswith("Assistant:"):
             text = text[len("Assistant:"):].strip()
-        return text
-
-
+        t_dec_end = perf_counter()
+
+        # Stats
+        input_tokens = int(inputs["input_ids"].shape[-1]) if "input_ids" in inputs else None
+        total_tokens = int(out_ids.shape[-1])  # includes prompt + generated
+        output_tokens = int(total_tokens - (input_tokens or 0)) if input_tokens is not None else None
+
+        pre_ms = (t_pre_end - t0) * 1000.0
+        infer_ms = (t_inf_end - t_inf_start) * 1000.0
+        decode_ms = (t_dec_end - t_inf_end) * 1000.0
+        total_ms = (t_dec_end - t0) * 1000.0
+
+        tps_infer = (output_tokens / ((t_inf_end - t_inf_start) + 1e-9)) if output_tokens else None
+        tps_total = (
+            (output_tokens / ((t_dec_end - t0) + 1e-9)) if output_tokens else None
+        )
+
+        gpu_mem_alloc_mb = gpu_mem_resv_mb = None
+        gpu_name = None
+        if self.device == "cuda":
+            try:
+                gpu_mem_alloc_mb = round(torch.cuda.max_memory_allocated() / (1024**2), 2)
+                gpu_mem_resv_mb = round(torch.cuda.max_memory_reserved() / (1024**2), 2)
+                gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
+            except Exception:
+                pass
+
+        metrics: Dict[str, Any] = {
+            **meta,
+            "gpu_name": gpu_name,
+            "timings_ms": {
+                "preprocess": round(pre_ms, 2),
+                "inference": round(infer_ms, 2),
+                "decode": round(decode_ms, 2),
+                "total": round(total_ms, 2),
+            },
+            "tokens": {
+                "input": input_tokens,
+                "output": output_tokens,
+                "total": total_tokens,
+            },
+            "throughput": {
+                "tokens_per_sec_inference": None if tps_infer is None else round(tps_infer, 2),
+                "tokens_per_sec_end_to_end": None if tps_total is None else round(tps_total, 2),
+            },
+            "gpu_memory_mb": {
+                "max_allocated": gpu_mem_alloc_mb,
+                "max_reserved": gpu_mem_resv_mb,
+            },
+        }
+
+        return (text, metrics) if return_stats else text
+
+
+# Convenience singleton
 _runner_singleton = None
 def get_runner():
     global _runner_singleton
@@ -95,4 +211,3 @@ def get_runner():
         _runner_singleton = SmolVLMRunner()
     return _runner_singleton
 
-