Z User
v3.1.0: Part B architecture migration — eliminate Hugging Face inference credit dependency
6c516b6 | 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() |