Upload folder using huggingface_hub

app/api/routes.py

@@ -15,6 +15,12 @@ from ..utils.cache import get_cache, set_cache
 from ..config import settings
 from ..db import repo
 
+# Strong references to in-flight background bulk jobs (the event loop only
+# keeps weak ones, so a task can vanish mid-run and its results are lost), and
+# a cap so a burst of bulk requests cannot spawn unbounded SMTP work.
+_background_tasks: set[asyncio.Task] = set()
+_bulk_semaphore = asyncio.Semaphore(settings.bulk_concurrency)
+
 router = APIRouter()
 
 
@@ -184,31 +190,47 @@ async def validate_bulk(request: BulkEmailRequest):
     task_id = uuid.uuid4().hex
     print(f"[BULK] Processing {len(request.emails)} emails in background (task {task_id})")
 
-    async def _run_and_store():
+    # Record the job BEFORE processing: without this row, finish_bulk_job is an
+    # UPDATE that matches nothing, and in-flight polls hit the "Redis down"
+    # branch and 404 — which the frontend treats as terminal.
+    if settings.enable_db_persistence:
         try:
-            results = await process_bulk_sync(request.emails, request.webhook_url, deep=request.deep)
-            if settings.enable_db_persistence:
-                try:
-                    repo.finish_bulk_job(task_id, results)
-                except Exception as exc:
-                    print(f"[BULK] finish_bulk_job failed: {exc}")
-            await set_cache(
-                f"bulk_result:{task_id}",
-                json.dumps({"status": "completed", "total": len(request.emails),
-                            "results": results}),
-                ttl=settings.redis_cache_ttl,
-            )
-            print(f"[BULK] task {task_id} completed: {len(results)} results")
+            repo.create_bulk_job(task_id, request.emails)
         except Exception as exc:
-            print(f"[BULK] task {task_id} failed: {exc}")
-            await set_cache(
-                f"bulk_result:{task_id}",
-                json.dumps({"status": "error", "error": str(exc),
-                            "results": []}),
-                ttl=settings.redis_cache_ttl,
-            )
+            print(f"[BULK] create_bulk_job failed: {exc}")
+
+    async def _run_and_store():
+        async with _bulk_semaphore:
+            try:
+                results = await process_bulk_sync(
+                    request.emails, request.webhook_url, deep=request.deep
+                )
+                if settings.enable_db_persistence:
+                    try:
+                        repo.finish_bulk_job(task_id, results)
+                    except Exception as exc:
+                        print(f"[BULK] finish_bulk_job failed: {exc}")
+                await set_cache(
+                    f"bulk_result:{task_id}",
+                    json.dumps({"status": "completed", "total": len(request.emails),
+                                "results": results}),
+                    ttl=settings.redis_cache_ttl,
+                )
+                print(f"[BULK] task {task_id} completed: {len(results)} results")
+            except Exception as exc:
+                print(f"[BULK] task {task_id} failed: {exc}")
+                await set_cache(
+                    f"bulk_result:{task_id}",
+                    json.dumps({"status": "error", "error": str(exc),
+                                "results": []}),
+                    ttl=settings.redis_cache_ttl,
+                )
 
-    asyncio.create_task(_run_and_store())
+    # The loop holds only a WEAK reference to tasks: without a strong ref the
+    # job can be garbage-collected mid-run and its results silently lost.
+    task = asyncio.create_task(_run_and_store())
+    _background_tasks.add(task)
+    task.add_done_callback(_background_tasks.discard)
 
     return JSONResponse(content={
         "task_id": task_id,
@@ -288,6 +310,11 @@ async def get_bulk_status(task_id: str):
                     total=job.get("total"),
                     results=_normalize_results(job.get("results"))
                 )
+        # An in-flight background job has no cached result yet. Report
+        # "processing" instead of 404 — the frontend treats 404 as terminal and
+        # would abandon a job that is still running.
+        if any(not t.done() for t in _background_tasks):
+            return TaskStatusResponse(task_id=task_id, status="processing")
         raise HTTPException(status_code=404, detail=f"Task not found: {task_id}")
 
     try:

app/main.py

@@ -68,6 +68,20 @@ async def lifespan(app: FastAPI):
     except Exception as exc:  # pragma: no cover - defensive
         print(f"[RETRY] drain loop skipped — {exc}")
 
+    # Warm the advisory model at startup: the first request otherwise pays the
+    # torch + checkpoint load (measured 42.6s cold start) even though inference
+    # now runs off the event loop. Fail-open and never fatal.
+    try:
+        from .ml.runtime import get_model_async
+
+        model = await get_model_async()
+        if model is not None and getattr(model, "loaded", False):
+            print("[ML] advisory model preloaded")
+        else:
+            print("[ML] advisory model unavailable - deterministic mode")
+    except Exception as exc:  # pragma: no cover - defensive
+        print(f"[ML] preload skipped — {exc}")
+
     yield
 
     if drain_task is not None:

app/ml/runtime.py

@@ -1,16 +1,24 @@
-"""Lean ONNX inference runtime for the email-risk model.
+"""Inference runtime for the email-risk model (PyTorch / transformers).
 
 Strictly optional: `load()` returning False (or `predict()` returning None)
-means "no ML this run" — the deterministic pipeline continues unchanged. Imports
-are lazy so the API boots with or without onnxruntime installed, and the test
-suite never requires torch.
-
-The deployed Space image must NOT contain torch (~2 GB); inference is ONNX only.
+means "no ML this run" — the deterministic pipeline continues unchanged.
+Imports are lazy so the API boots with or without torch installed, and the test
+suite never requires a model download.
+
+The deployed Space image DOES contain CPU torch (Dockerfile installs it from
+the PyTorch CPU index) and ships the bundled model at `app/ml/model`. Do not
+strip either: the loader fails open, so removing them silently disables the
+advisory layer instead of raising.
+
+CPU inference is offloaded to a worker thread (`predict_async`, `load_async`)
+because running it inline stalls every concurrent request — measured at 42.6s
+cold start before this.
 """
 from __future__ import annotations
 
+import asyncio
 import threading
-from typing import Optional
+from typing import Callable, Optional
 
 from ..config import settings
 
@@ -82,6 +90,15 @@ class EmailRiskModel:
                 self.last_error = f"{type(exc).__name__}: {exc}"
                 return False
 
+    async def load_async(self) -> bool:
+        """Load the model without blocking the event loop.
+
+        Loading torch + a checkpoint takes seconds; doing it inline in an async
+        handler stalls every concurrent request (measured: 42.6s cold start
+        before this).
+        """
+        return await asyncio.to_thread(self.load)
+
     def predict(self, texts: list[str]) -> Optional[list[tuple[str, float]]]:
         """Classify texts. Returns (LABEL, confidence) tuples, or None on failure.
 
@@ -93,29 +110,61 @@ class EmailRiskModel:
             return None
         try:
             preds = self._pipeline(texts, batch_size=self.batch_size)
-        except Exception:
+            out: list[tuple[str, float]] = []
+            for batch in preds:
+                # transformers may return a flat list or one list per element.
+                if isinstance(batch, list):
+                    preds_nested = batch
+                else:
+                    preds_nested = [batch]
+                for p in preds_nested:
+                    # An unexpected payload shape must fail open, not raise:
+                    # the batch path calls predict() outside its own try.
+                    if not isinstance(p, dict):
+                        raise TypeError(f"unexpected prediction: {type(p).__name__}")
+                    label = str(p.get("label", "")).upper()
+                    try:
+                        score = float(p.get("score", 0.0))
+                    except (TypeError, ValueError):
+                        score = 0.0
+                    out.append((label, score))
+            return out if len(out) == len(texts) else None
+        except Exception as exc:
+            self.last_error = f"predict: {type(exc).__name__}: {exc}"
             return None
-        out: list[tuple[str, float]] = []
-        for batch in preds:
-            # transformers may return a flat list or one list per batch element.
-            if isinstance(batch, list):
-                preds_nested = batch
-            else:
-                preds_nested = [batch]
-            for p in preds_nested:
-                label = str(p.get("label", "")).upper()
-                try:
-                    score = float(p.get("score", 0.0))
-                except (TypeError, ValueError):
-                    score = 0.0
-                out.append((label, score))
-        return out if len(out) == len(texts) else None
 
 
 _model: Optional[EmailRiskModel] = None
 _model_lock = threading.Lock()
 
 
+async def predict_async(model, texts: list[str]):
+    """Classify texts without blocking the event loop.
+
+    Works with either an `EmailRiskModel` (CPU inference in a worker thread) or
+    any injected test double exposing a sync `predict` (called directly).
+    """
+    if model is None or not texts:
+        return None
+    predict = getattr(model, "predict", None)
+    if predict is None:
+        return None
+    if isinstance(model, EmailRiskModel):
+        return await asyncio.to_thread(predict, texts)
+    return predict(texts)
+
+
+async def get_model_async() -> Optional[EmailRiskModel]:
+    """`get_model()` with the (slow) first load kept off the event loop."""
+    global _model
+    if not settings.ml_disposable_enabled or not settings.ml_model_repo:
+        return None
+    if _model is not None:
+        return _model
+    await asyncio.to_thread(get_model)
+    return _model
+
+
 def get_model() -> Optional[EmailRiskModel]:
     """Process-wide model, built AND loaded when flag + repo are configured."""
     global _model

app/validator/engine.py

@@ -20,7 +20,7 @@ from .syntax import validate_syntax
 from .dns_check import check_dns
 from .disposable import check_disposable, domain_in_blocklist, domain_is_known
 from .mailbox_provider import MailboxProvider, resolve_provider
-from ..ml.runtime import get_model
+from ..ml.runtime import get_model_async, predict_async
 from .mx_intel import mx_is_trusted
 from .retry_queue import is_greylist_result
 from .scoring import calculate_score
@@ -233,7 +233,7 @@ class EmailValidationEngine:
         # learned 'unfamiliar domain string = disposable' from an imbalanced
         # dataset and flags legitimate small businesses on Google Workspace /
         # M365 — the deterministic MX intelligence is the better signal there.
-        ml_model = get_model()
+        ml_model = await get_model_async()
         mx_host = (
             dns_result.mx_records[0]["host"]
             if dns_result.mx_records
@@ -245,10 +245,9 @@ class EmailValidationEngine:
             and deep
             and not mx_is_trusted(mx_host)
         ):
-            try:
-                preds = ml_model.predict([f"john.smith [SEP] {domain}"])
-            except Exception:
-                preds = None
+            # Off the event loop: CPU inference inline would stall every other
+            # concurrent request (measured 42.6s cold start before this).
+            preds = await predict_async(ml_model, [f"john.smith [SEP] {domain}"])
             if preds:
                 disposable = disposable.model_copy(
                     update={"ml_label": preds[0][0], "ml_confidence": preds[0][1]}
@@ -546,7 +545,7 @@ class EmailValidationEngine:
             # there, and the model's 'unfamiliar string' bias would otherwise
             # flag legitimate businesses as disposable.
             ml_map: dict[str, tuple[str, float]] = {}
-            ml_model = get_model()
+            ml_model = await get_model_async()
             if ml_model is not None:
                 unknown_domains = sorted(
                     {
@@ -558,8 +557,11 @@ class EmailValidationEngine:
                     }
                 )
                 if unknown_domains:
-                    preds = ml_model.predict(
-                        [f"john.smith [SEP] {d}" for d in unknown_domains]
+                    # One batched call, off the event loop (CPU inference inline
+                    # stalls concurrent requests), and fail-open by construction:
+                    # predict_async never raises.
+                    preds = await predict_async(
+                        ml_model, [f"john.smith [SEP] {d}" for d in unknown_domains]
                     )
                     if preds:
                         ml_map = dict(zip(unknown_domains, preds))

tests/test_engine_ml.py

@@ -54,7 +54,8 @@ def _engine_with_ml(flag=True, repo="user/model", predict=None):
         fake_model.predict = MagicMock(side_effect=predict)
         fake_model.load = MagicMock(return_value=True)
         engine.mailbox = engine.mailbox  # unchanged
-        patcher = patch("app.validator.engine.get_model", return_value=fake_model)
+        patcher = patch("app.validator.engine.get_model_async",
+                        new=AsyncMock(return_value=fake_model))
         patcher.start()
         engine._ml_patcher = patcher  # keep alive until stopAll
         engine._ml_model = fake_model

tests/test_ml_async.py

@@ -0,0 +1,101 @@
+import asyncio
+import time
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from app.ml.runtime import EmailRiskModel, predict_async
+
+
+def _slow_pipeline(texts, batch_size=None):
+    time.sleep(0.4)
+    return [[{"label": "NORMAL", "score": 0.9}] for _ in texts]
+
+
+class NonBlockingInferenceTests(unittest.IsolatedAsyncioTestCase):
+    """CPU inference must run OFF the event loop.
+
+    Measured in production: the first request after a cold start took 42.6s and
+    warm requests ~2s because torch inference ran synchronously inside `async
+    def` handlers, stalling every concurrent request (including /health).
+    """
+
+    async def test_predict_async_does_not_block_the_event_loop(self):
+        m = EmailRiskModel(repo="user/model")
+        m._pipeline = _slow_pipeline
+        m._loaded = True
+
+        ticks = 0
+
+        async def ticker():
+            nonlocal ticks
+            while True:
+                await asyncio.sleep(0.02)
+                ticks += 1
+
+        t = asyncio.create_task(ticker())
+        out = await predict_async(m, ["a [SEP] b.test"])
+        t.cancel()
+
+        self.assertEqual(out, [("NORMAL", 0.9)])
+        # A blocking implementation would leave the loop frozen for ~0.4s
+        # (ticks == 0); off-loop it keeps ticking throughout.
+        self.assertGreaterEqual(
+            ticks, 5, f"event loop was blocked during inference (ticks={ticks})"
+        )
+
+    async def test_load_async_runs_the_factory_off_the_event_loop(self):
+        def slow_factory(repo, token, batch_size):
+            time.sleep(0.3)
+            return _slow_pipeline
+
+        m = EmailRiskModel(repo="user/model", pipeline_factory=slow_factory)
+        ticks = 0
+
+        async def ticker():
+            nonlocal ticks
+            while True:
+                await asyncio.sleep(0.02)
+                ticks += 1
+
+        t = asyncio.create_task(ticker())
+        self.assertTrue(await m.load_async())
+        t.cancel()
+        self.assertGreaterEqual(
+            ticks, 4, f"model load blocked the event loop (ticks={ticks})"
+        )
+
+    async def test_predict_async_is_fail_open(self):
+        m = EmailRiskModel(repo="user/model")
+        m._pipeline = MagicMock(side_effect=RuntimeError("boom"))
+        m._loaded = True
+        self.assertIsNone(await predict_async(m, ["a [SEP] b.test"]))
+
+    async def test_predict_async_handles_unloaded_model(self):
+        m = EmailRiskModel(repo="user/model")
+        self.assertIsNone(await predict_async(m, ["a [SEP] b.test"]))
+
+    async def test_predict_async_matches_sync_results_for_fakes(self):
+        # Engine tests inject MagicMock models with a sync predict; the async
+        # helper must work with those unchanged.
+        fake = MagicMock()
+        fake.predict = MagicMock(return_value=[("DISPOSABLE", 0.97)])
+        out = await predict_async(fake, ["x [SEP] y.test"])
+        self.assertEqual(out, [("DISPOSABLE", 0.97)])
+        fake.predict.assert_called_once()
+
+    async def test_concurrent_predictions_are_parallelised(self):
+        """Two 0.3s inferences should overlap, not serialise (thread pool)."""
+        m = EmailRiskModel(repo="user/model")
+        m._pipeline = _slow_pipeline
+        m._loaded = True
+        t0 = time.perf_counter()
+        await asyncio.gather(
+            predict_async(m, ["a [SEP] b.test"]),
+            predict_async(m, ["c [SEP] d.test"]),
+        )
+        elapsed = time.perf_counter() - t0
+        self.assertLess(elapsed, 0.55, f"inferences serialised ({elapsed:.2f}s)")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_ml_mx_gate.py

@@ -61,8 +61,8 @@ class TrustedMxGateTests(unittest.IsolatedAsyncioTestCase):
 
     def setUp(self):
         self.patchers = [
-            patch("app.validator.engine.get_model",
-                  return_value=_fake_model("DISPOSABLE")),
+            patch("app.validator.engine.get_model_async",
+                  new=AsyncMock(return_value=_fake_model("DISPOSABLE"))),
         ]
         for p in self.patchers:
             p.start()