fix(intern-decision-serve): load and score on one dedicated inference thread
torch keeps CUDA state per host thread (cuBLAS handles and workspaces), partly outside the per-process VRAM cap. Scoring on anyio's threadpool let 40 threads each create it: measured on fv-ml1 GPU 3, +252 MiB outside the cap and +326 MiB inside, which pushed the process past the 10,300 MiB GPU 1 budget. Load, warm-up and every call now run on the same single thread.
This commit is contained in:
@@ -432,3 +432,35 @@ def test_more_than_max_queue_requests_in_progress_get_429_busy_before_the_body_i
|
||||
engine.release.set()
|
||||
assert [f.result().status_code for f in held] == [200, 200]
|
||||
assert make_client(FakeEngine(), max_queue=2).post("/decide", json=ROW, headers=AUTH).status_code == 200
|
||||
|
||||
|
||||
class ThreadRecordingEngine(FakeEngine):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.threads = set()
|
||||
|
||||
def predict(self, request):
|
||||
self.threads.add(threading.get_ident())
|
||||
return super().predict(request)
|
||||
|
||||
|
||||
def test_every_call_runs_on_one_dedicated_inference_thread():
|
||||
"""torch keeps CUDA state per host thread (cuBLAS handles + workspaces), outside the VRAM cap;
|
||||
measured 2026-09-30: 40 worker threads added 252 MiB outside the cap. So: one thread, always."""
|
||||
engine = ThreadRecordingEngine()
|
||||
with make_client(engine) as client, ThreadPoolExecutor(12) as pool:
|
||||
futures = [pool.submit(client.post, "/decide/shared", json={"state": f"s{i}", "decisions": decisions(20)},
|
||||
headers=AUTH) for i in range(24)]
|
||||
assert all(f.result().status_code == 200 for f in futures)
|
||||
assert len(engine.threads) == 1
|
||||
|
||||
|
||||
def test_the_app_uses_the_executor_it_is_given():
|
||||
from intern_decision_serve.app import create_app as build
|
||||
engine = ThreadRecordingEngine()
|
||||
executor = ThreadPoolExecutor(1)
|
||||
home = executor.submit(threading.get_ident).result()
|
||||
client = TestClient(build(Settings(api_token=TOKEN), engine, executor=executor))
|
||||
assert client.post("/decide", json=ROW, headers=AUTH).status_code == 200
|
||||
assert engine.threads == {home}
|
||||
executor.shutdown()
|
||||
|
||||
@@ -21,3 +21,28 @@ def test_app_from_env_goes_offline_before_the_engine_loads(monkeypatch):
|
||||
app = main.app_from_env()
|
||||
assert seen == {"offline": ("1", "1"), "token": "k" * 40}
|
||||
assert app.title == "intern-decision-serve"
|
||||
|
||||
|
||||
def test_the_engine_loads_on_the_same_thread_every_call_later_runs_on(monkeypatch):
|
||||
import threading
|
||||
from intern_decision_serve import engine as engine_module
|
||||
from intern_decision_serve import main
|
||||
from fastapi.testclient import TestClient
|
||||
seen = {}
|
||||
|
||||
class Recording(FakeEngine):
|
||||
def predict(self, request):
|
||||
seen.setdefault("calls", set()).add(threading.get_ident())
|
||||
return super().predict(request)
|
||||
|
||||
def fake_load(settings, **_kw):
|
||||
seen["load"] = threading.get_ident()
|
||||
return Recording()
|
||||
|
||||
monkeypatch.setenv("INTERN_DECISION_API_TOKEN", "k" * 40)
|
||||
monkeypatch.setattr(engine_module.TorchEngine, "load", staticmethod(fake_load))
|
||||
client = TestClient(main.app_from_env())
|
||||
body = {"id": "r", "state": "s", "question": "q?", "options": [{"id": "a", "description": "A"},
|
||||
{"id": "b", "description": "B"}]}
|
||||
assert client.post("/decide", json=body, headers={"Authorization": "Bearer " + "k" * 40}).status_code == 200
|
||||
assert seen["calls"] == {seen["load"]} and seen["load"] != threading.get_ident()
|
||||
|
||||
Reference in New Issue
Block a user