Sync main.py from blood-brain-omics repo

main.py

@@ -27,6 +27,7 @@ Endpoints:
     GET /api/v1/ready                 - Readiness probe (SELECT 1)
 """
 
+import asyncio
 import json
 import logging
 import os
@@ -38,7 +39,7 @@ from typing import Optional
 import duckdb
 from fastapi import FastAPI, HTTPException, Query, Request
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, Response
 
 logging.basicConfig(
     level=logging.INFO,
@@ -68,6 +69,14 @@ DUCKDB_MEMORY_LIMIT = os.environ.get("DUCKDB_MEMORY_LIMIT", "512MB")
 DUCKDB_THREADS = int(os.environ.get("DUCKDB_THREADS", "1"))
 DUCKDB_TEMP_DIR = os.environ.get("DUCKDB_TEMP_DIR", "/tmp/duckdb_tmp")
 
+# Concurrency cap — bound in-flight queries to avoid OOM/process crash on
+# CPU Basic. Excess requests get 429 and the frontend can retry.
+MAX_INFLIGHT = int(os.environ.get("API_MAX_INFLIGHT", "8"))
+
+# Response cache TTL (seconds) — absorbs repeat queries from rapid clicks.
+CACHE_TTL = float(os.environ.get("API_CACHE_TTL", "60"))
+CACHE_MAX_ENTRIES = int(os.environ.get("API_CACHE_MAX_ENTRIES", "256"))
+
 app = FastAPI(
     title="Blood-Brain Omics Benchmark API",
     version="1.0.0",
@@ -82,6 +91,99 @@ app.add_middleware(
 )
 
 
+# =============================================================================
+# Concurrency cap + response cache
+# =============================================================================
+
+# Bound the number of in-flight requests so that bursts cannot OOM/crash the
+# DuckDB process on CPU Basic. Excess requests fast-fail with 429 instead of
+# wedging the worker.
+_inflight_sem = asyncio.Semaphore(MAX_INFLIGHT)
+
+# Tiny TTL cache keyed on (path, query_string). Frontend rapid clicks
+# (e.g. switching panels) repeatedly hit the same URLs — caching them
+# absorbs the load before it ever reaches DuckDB.
+_cache: dict[str, tuple[float, int, bytes, str]] = {}
+_CACHEABLE_PATH_PREFIXES = (
+    "/api/v1/registry",
+    "/api/v1/maxn/",
+    "/api/v1/h2h/",
+    "/api/v1/temporal",
+    "/api/v1/features",
+)
+
+
+def _cache_key(request: Request) -> Optional[str]:
+    if request.method != "GET":
+        return None
+    path = request.url.path
+    if not any(path.startswith(p) for p in _CACHEABLE_PATH_PREFIXES):
+        return None
+    return f"{path}?{request.url.query}"
+
+
+def _cache_get(key: str) -> Optional[tuple[int, bytes, str]]:
+    entry = _cache.get(key)
+    if entry is None:
+        return None
+    expires, status, body, ctype = entry
+    if expires < time.monotonic():
+        _cache.pop(key, None)
+        return None
+    return status, body, ctype
+
+
+def _cache_put(key: str, status: int, body: bytes, ctype: str) -> None:
+    if status != 200:
+        return
+    if len(_cache) >= CACHE_MAX_ENTRIES:
+        # Cheap eviction: drop the oldest-expiring entry.
+        oldest = min(_cache.items(), key=lambda kv: kv[1][0])[0]
+        _cache.pop(oldest, None)
+    _cache[key] = (time.monotonic() + CACHE_TTL, status, body, ctype)
+
+
+@app.middleware("http")
+async def cache_and_throttle(request: Request, call_next):
+    key = _cache_key(request)
+    if key is not None:
+        hit = _cache_get(key)
+        if hit is not None:
+            status, body, ctype = hit
+            return Response(
+                content=body, status_code=status, media_type=ctype,
+                headers={"x-cache": "hit"},
+            )
+
+    # Bound concurrency. If saturated, fail fast with 429 rather than queue
+    # indefinitely (queueing under sustained load is what crashes the worker).
+    try:
+        await asyncio.wait_for(_inflight_sem.acquire(), timeout=0.05)
+    except asyncio.TimeoutError:
+        return JSONResponse(
+            status_code=429,
+            content={"error": "busy", "detail": "Server busy, please retry."},
+            headers={"retry-after": "1"},
+        )
+
+    try:
+        response = await call_next(request)
+    finally:
+        _inflight_sem.release()
+
+    if key is not None and response.status_code == 200:
+        # Buffer the streaming body so we can both cache and return it.
+        body_chunks = [chunk async for chunk in response.body_iterator]
+        body = b"".join(body_chunks)
+        ctype = response.headers.get("content-type", "application/json")
+        _cache_put(key, response.status_code, body, ctype)
+        return Response(
+            content=body, status_code=response.status_code, media_type=ctype,
+            headers={"x-cache": "miss"},
+        )
+    return response
+
+
 # =============================================================================
 # Middleware: request logging + unhandled-exception handler
 # =============================================================================
@@ -212,6 +314,19 @@ def shutdown():
 VALID_METRICS = {"r2", "pearson", "mse", "mae"}
 
 
+def q(sql: str, params=None):
+    """Run a query on an isolated DuckDB cursor.
+
+    DuckDB Connection objects are NOT safe for concurrent use across
+    requests — sharing the global ``db`` connection between overlapping
+    FastAPI requests causes corrupted result state and hung responses.
+    Cursors created via ``db.cursor()`` are cheap, isolated, and safe
+    for concurrent queries against the same in-memory database.
+    """
+    cur = db.cursor()
+    return cur.execute(sql, params) if params is not None else cur.execute(sql)
+
+
 def validate_metric(metric: str) -> str:
     if metric not in VALID_METRICS:
         raise HTTPException(400, f"Invalid metric: {metric}. "
@@ -253,7 +368,7 @@ def maxn_heatmap(
     validate_metric(metric)
     phase = covariate_phase("maxn", covariates)
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT blood_platform, brain_target,
                MEDIAN({metric}) as value,
                MEDIAN(n_samples) as n_samples,
@@ -305,7 +420,7 @@ def maxn_detail(
     """Per-module results for one blood x brain combination."""
     phase = covariate_phase("maxn", covariates)
 
-    rows = db.execute("""
+    rows = q("""
         SELECT target, n_samples, r2, pearson, mse, mae
         FROM results
         WHERE phase = ? AND blood_platform = ? AND brain_target = ?
@@ -338,7 +453,7 @@ def h2h_blood(
     validate_metric(metric)
     phase = "h2h_blood_withcov" if covariates == "included" else "h2h_blood"
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT target, blood_platform, {metric}
         FROM results
         WHERE phase = ? AND h2h_pair = ? AND brain_target = ? AND model = ?
@@ -375,7 +490,7 @@ def h2h_blood_summary(
     validate_metric(metric)
     phase = "h2h_blood_withcov" if covariates == "included" else "h2h_blood"
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT h2h_pair, brain_target, target, blood_platform, {metric}
         FROM results
         WHERE phase = ? AND model = ? AND h2h_pair IS NOT NULL
@@ -434,7 +549,7 @@ def h2h_brain(
     validate_metric(metric)
     phase = "h2h_brain_withcov" if covariates == "included" else "h2h_brain"
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT h2h_pair, brain_target, target, {metric}
         FROM results
         WHERE phase = ? AND blood_platform = ? AND model = ?
@@ -468,7 +583,7 @@ def h2h_models(
     validate_metric(metric)
     phase = covariate_phase("maxn", covariates)
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT target, model, {metric}
         FROM results
         WHERE phase = ? AND blood_platform = ? AND brain_target = ?
@@ -506,14 +621,14 @@ def temporal(
     phase = "temporal_withcov" if covariates == "included" else "temporal"
 
     if target:
-        rows = db.execute(f"""
+        rows = q(f"""
             SELECT temporal_bin, brain_target, target, {metric}, n_samples
             FROM results
             WHERE phase = ? AND model = ? AND target = ?
             ORDER BY temporal_bin
         """, [phase, model, target]).fetchall()
     else:
-        rows = db.execute(f"""
+        rows = q(f"""
             SELECT temporal_bin, brain_target, target, {metric}, n_samples
             FROM results
             WHERE phase = ? AND model = ?
@@ -544,7 +659,7 @@ def feature_importance(
     limit: int = Query(30, ge=1, le=100),
 ):
     """Top features for a specific module/target."""
-    rows = db.execute("""
+    rows = q("""
         SELECT feature_name, importance, rank
         FROM features
         WHERE phase = ? AND blood_platform = ? AND brain_target = ?
@@ -574,7 +689,7 @@ def feature_cross_target(
     limit: int = Query(30, ge=1, le=100),
 ):
     """Aggregated feature importance across all modules in a brain target."""
-    rows = db.execute("""
+    rows = q("""
         SELECT feature_name,
                AVG(importance) as mean_importance,
                COUNT(DISTINCT target) as n_targets,
@@ -613,7 +728,7 @@ def maxn_panel_summary(
     validate_metric(metric)
     phase = covariate_phase("maxn", covariates)
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT brain_target,
                AVG({metric}) as mean_val,
                MIN({metric}) as min_val,
@@ -659,7 +774,7 @@ def maxn_outcome_summary(
     validate_metric(metric)
     phase = covariate_phase("maxn", covariates)
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT blood_platform,
                AVG({metric}) as mean_val,
                MIN({metric}) as min_val,
@@ -706,7 +821,7 @@ def maxn_target_comparison(
     validate_metric(metric)
     phase = covariate_phase("maxn", covariates)
 
-    rows = db.execute(f"""
+    rows = q(f"""
         SELECT blood_platform, {metric}, n_samples
         FROM results
         WHERE phase = ? AND brain_target = ? AND target = ? AND model = ?
@@ -738,8 +853,8 @@ def health():
             content={"status": "degraded", "detail": "Database not initialized"},
         )
     try:
-        n_results = db.execute("SELECT COUNT(*) FROM results").fetchone()[0]
-        n_features = db.execute("SELECT COUNT(*) FROM features").fetchone()[0]
+        n_results = q("SELECT COUNT(*) FROM results").fetchone()[0]
+        n_features = q("SELECT COUNT(*) FROM features").fetchone()[0]
     except Exception as exc:
         logger.error("Health check query failed: %s", exc)
         return JSONResponse(
@@ -758,7 +873,7 @@ def ready():
             content={"ready": False, "detail": "Database not initialized"},
         )
     try:
-        db.execute("SELECT 1").fetchone()
+        q("SELECT 1").fetchone()
     except Exception as exc:
         logger.error("Ready check query failed: %s", exc)
         return JSONResponse(