Update app.py

app.py

@@ -1,17 +1,38 @@
 """
-CLOAK PII Demo — ONNX INT8 CPU inference + Gradio UI + REST /detect API.
-Downloads model.int8.onnx from HF Hub at startup.
-Set HF_TOKEN secret in Space settings if the repo is private.
-The /detect endpoint is what the CLOAK API (Next.js) calls via callHfSpace:
-  POST /detect  Authorization: Bearer <API_SECRET>  {"text": "..."}
-  -> {entities, nested, redacted_text, model, processing_time_ms}
+CLOAK PII API — ONNX INT8 CPU inference + Gradio demo UI, tuned for the CLOAK
+DLP server (extension analyze, email/chat/drive DLP, and the Drive DSPM crawl).
+
+Endpoints (FastAPI, mounted alongside the Gradio UI):
+  GET  /health                          -> {status, model, engine, threads, classes}
+  POST /detect   Bearer <API_SECRET>    {"text": "..."}
+       -> {entities, nested, redacted_text, model, processing_time_ms}
+
+Contract notes for the CLOAK server (server/src/lib/pii.ts):
+  * entities: [{type, text, start, end, confidence}] — offsets into the sent text.
+  * 401 on a bad key, 503 on inference failure, 422 on oversize body — the
+    server records these as scan ERRORS, never as "clean". Do not return
+    200-with-empty on failure paths.
+  * MAX_TEXT_CHARS here must match DETECT_MAX_CHARS in pii.ts (30_000).
+
+Space secrets:
+  HF_TOKEN    — read access to the private model repo (also the default secret).
+  API_SECRET  — Bearer token CLOAK sends; defaults to HF_TOKEN when unset.
+  MODEL_REPO  — override the model repo id (default Wardline/CLOAK_V3.6).
 """
 
+import hmac
 import html as _html
 import json
 import os
+import threading
 import time
 
+# HF's Gradio-Space image sets GRADIO_SSR_MODE=true, which spawns a second
+# (Node SSR) server on port 7861 — it double-binds under the
+# mount_gradio_app + uvicorn pattern and crashes the Space (exit 3).
+# Must be forced off BEFORE gradio is imported.
+os.environ["GRADIO_SSR_MODE"] = "false"
+
 import gradio as gr
 import numpy as np
 import onnxruntime as ort
@@ -25,29 +46,27 @@ from transformers import AutoTokenizer
 
 MODEL_REPO = os.getenv("MODEL_REPO", "Wardline/CLOAK_V3.6")
 HF_TOKEN   = os.getenv("HF_TOKEN")
-# Bearer token the CLOAK API must send. callHfSpace sends `Bearer ${HF_TOKEN}`,
-# so default API_SECRET to HF_TOKEN — the same secret works for both unless set.
 API_SECRET = os.getenv("API_SECRET", HF_TOKEN)
 
-MAX_LEN   = 512
-STRIDE    = 128
-THRESHOLD = 0.0
+MAX_LEN        = 512      # model window (tokens)
+STRIDE         = 128      # overlap between windows — no blind spots at joins
+THRESHOLD      = 0.0      # logit threshold (UI slider can tighten it)
+MAX_TEXT_CHARS = 30_000   # keep in sync with DETECT_MAX_CHARS in pii.ts
+MAX_WINDOWS    = 96       # hard ceiling on windows per request (runaway guard)
 
-# ── Thread count: read cgroup CPU quota, not sched_getaffinity ────────────────
-# sched_getaffinity returns the host's full CPU set inside Docker containers.
-# The real allocated vCPU count lives in the cgroup CFS quota.
+# ── Thread count: cgroup CPU quota, not sched_getaffinity ─────────────────────
+# sched_getaffinity reports the HOST's cores inside a container; the real vCPU
+# allocation lives in the cgroup CFS quota.
 
 def _container_cpu_count() -> int:
-    # cgroup v2
-    try:
+    try:  # cgroup v2
         with open("/sys/fs/cgroup/cpu.max") as f:
             quota, period = f.read().split()
         if quota != "max":
             return max(1, round(int(quota) / int(period)))
     except Exception:
         pass
-    # cgroup v1
-    try:
+    try:  # cgroup v1
         with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as f:
             quota = int(f.read())
         with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as f:
@@ -58,6 +77,7 @@ def _container_cpu_count() -> int:
         pass
     return os.cpu_count() or 2
 
+
 _n_threads = _container_cpu_count()
 print(f"Container CPU count: {_n_threads}  (os.cpu_count={os.cpu_count()})")
 
@@ -81,6 +101,11 @@ _opts.inter_op_num_threads = 1
 _opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
 _sess = ort.InferenceSession(_int8_path, sess_options=_opts, providers=["CPUExecutionProvider"])
 
+# One inference at a time. ONNX already saturates every allotted core per run —
+# concurrent runs would thrash the CPU and slow BOTH callers down. Waiters queue
+# here (FastAPI runs sync endpoints in a threadpool, so requests just block).
+_infer_lock = threading.Lock()
+
 print(f"ONNX INT8 session ready — {_n_threads} threads, {len(_id2label)} classes")
 
 
@@ -104,31 +129,34 @@ def detect(text: str, threshold: float = THRESHOLD):
     all_offsets = enc.pop("offset_mapping")
     enc.pop("overflow_to_sample_mapping", None)
 
+    n_windows = min(all_offsets.shape[0], MAX_WINDOWS)
+
     best: dict = {}
-    for w in range(all_offsets.shape[0]):
-        feed = {
-            "input_ids":      enc["input_ids"][w:w + 1].astype(np.int64),
-            "attention_mask": enc["attention_mask"][w:w + 1].astype(np.int64),
-        }
-        logits  = _sess.run(["logits"], feed)[0][0]   # [classes, L, L]
-        offsets = all_offsets[w]
-
-        cidx, sidx, eidx = np.where(logits > threshold)
-        if len(cidx) == 0:
-            continue
+    with _infer_lock:
+        for w in range(n_windows):
+            feed = {
+                "input_ids":      enc["input_ids"][w:w + 1].astype(np.int64),
+                "attention_mask": enc["attention_mask"][w:w + 1].astype(np.int64),
+            }
+            logits  = _sess.run(["logits"], feed)[0][0]   # [classes, L, L]
+            offsets = all_offsets[w]
+
+            cidx, sidx, eidx = np.where(logits > threshold)
+            if len(cidx) == 0:
+                continue
 
-        scores = 1.0 / (1.0 + np.exp(-logits[cidx, sidx, eidx]))
-        cs = offsets[sidx, 0]
-        ce = offsets[eidx, 1]
-        keep = cs < ce
-        for cid, c0, c1, sc in zip(cidx[keep], cs[keep], ce[keep], scores[keep]):
-            key = (int(cid), int(c0), int(c1))
-            if key not in best or sc > best[key]["score"]:
-                best[key] = {
-                    "entity_group": _id2label[int(cid)],
-                    "start": int(c0), "end": int(c1),
-                    "score": float(sc),
-                }
+            scores = 1.0 / (1.0 + np.exp(-logits[cidx, sidx, eidx]))
+            cs = offsets[sidx, 0]
+            ce = offsets[eidx, 1]
+            keep = cs < ce
+            for cid, c0, c1, sc in zip(cidx[keep], cs[keep], ce[keep], scores[keep]):
+                key = (int(cid), int(c0), int(c1))
+                if key not in best or sc > best[key]["score"]:
+                    best[key] = {
+                        "entity_group": _id2label[int(cid)],
+                        "start": int(c0), "end": int(c1),
+                        "score": float(sc),
+                    }
 
     doc_len = max(1, len(text))
     cands = [e for e in best.values() if (e["end"] - e["start"]) <= 0.5 * doc_len]
@@ -158,6 +186,7 @@ COLOR_MAP = {
 _FALLBACK = ["#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD", "#F7DC6F", "#85C1E9"]
 _extra: dict = {}
 
+
 def _label_color(label: str) -> str:
     if label in COLOR_MAP:
         return COLOR_MAP[label]
@@ -213,6 +242,7 @@ def run_analyze(text: str, threshold: float):
     if not text or not text.strip():
         return "<div style='color:#888'>Paste text above and click Analyze.</div>", pd.DataFrame(), "", ""
 
+    text = text[:MAX_TEXT_CHARS]
     t0 = time.time()
     flat, cands = detect(text, threshold)
     elapsed = round(time.time() - t0, 2)
@@ -296,17 +326,21 @@ with gr.Blocks(title="CLOAK PII Detector", theme=gr.themes.Soft()) as demo:
         outputs=[txt_input, entity_table, redacted_out, summary_box],
     )
 
-# ── REST API for the CLOAK backend (callHfSpace contract) ─────────────────────
+# Cap simultaneous UI jobs so demo users can't starve the REST API.
+demo.queue(default_concurrency_limit=1, max_size=8)
 
-api = FastAPI(title="CLOAK PII API", version="1.0.0")
+
+# ── REST API for the CLOAK backend ────────────────────────────────────────────
+
+api = FastAPI(title="CLOAK PII API", version="1.1.0")
 
 
 class DetectRequest(BaseModel):
-    text: str = Field(..., min_length=1, max_length=20_000)
+    text: str = Field(..., min_length=1, max_length=MAX_TEXT_CHARS)
 
 
 def _fmt(text: str, e: dict) -> dict:
-    """Shape one entity to the contract the CLOAK API expects."""
+    """Shape one entity to the contract the CLOAK server expects."""
     return {
         "type":       e["entity_group"],
         "text":       text[e["start"]:e["end"]].strip(),
@@ -321,14 +355,16 @@ def _verify(request: Request):
     if not auth.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing Authorization header")
     token = auth.removeprefix("Bearer ").strip()
-    if not API_SECRET or token != API_SECRET:
+    # Constant-time compare — a plain != leaks timing.
+    if not API_SECRET or not hmac.compare_digest(token, API_SECRET):
         raise HTTPException(status_code=401, detail="Invalid token")
 
 
 @api.get("/health")
 def health():
     return {"status": "ok", "model": MODEL_REPO, "engine": "onnx-int8-cpu",
-            "threads": _n_threads, "classes": len(_id2label)}
+            "threads": _n_threads, "classes": len(_id2label),
+            "max_chars": MAX_TEXT_CHARS}
 
 
 @api.post("/detect")
@@ -338,6 +374,8 @@ def detect_api(body: DetectRequest, request: Request):
     try:
         flat, cands = detect(body.text)
     except Exception as exc:
+        # 503, never 200-with-empty: the CLOAK server must record scan errors,
+        # not false "clean" results.
         raise HTTPException(status_code=503, detail=f"Inference failed: {exc}")
 
     entities = [_fmt(body.text, e) for e in flat
@@ -350,20 +388,23 @@ def detect_api(body: DetectRequest, request: Request):
     for e in sorted(flat, key=lambda x: x["start"], reverse=True):
         redacted = redacted[:e["start"]] + f'[{e["entity_group"]}]' + redacted[e["end"]:]
 
+    ms = int((time.time() - t0) * 1000)
+    print(f"[detect] {len(body.text)} chars -> {len(entities)} entities in {ms}ms")
+
     return {
         "entities":           entities,
         "nested":             nested,
         "redacted_text":      redacted,
         "model":              MODEL_REPO,
-        "processing_time_ms": int((time.time() - t0) * 1000),
+        "processing_time_ms": ms,
     }
 
 
-# Mount the Gradio UI at "/" on the same FastAPI app, so one Space serves both
+# Mount the Gradio UI at "/" on the same FastAPI app — one Space serves both
 # the demo UI and the /detect + /health API.
 app = gr.mount_gradio_app(api, demo, path="/")
 
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))