File size: 3,782 Bytes
d82b7e3 6c516b6 d82b7e3 d92dd26 d82b7e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | 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, 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
# Serial execution would take ~0.8 s (two full sleeps); anything
# comfortably below that proves overlap. 0.7 leaves scheduling slack
# on a loaded machine while still failing on serialisation.
self.assertLess(elapsed, 0.7, f"inferences serialised ({elapsed:.2f}s)")
if __name__ == "__main__":
unittest.main() |