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:
vh
2026-09-30 09:18:46 -07:00
parent f7415db5c9
commit f21369e4ac
5 changed files with 87 additions and 11 deletions
@@ -6,17 +6,18 @@ in calls of at most 16, and the answers are re-keyed into semif's result shape.
"""
from __future__ import annotations
import asyncio
import hmac
import itertools
import json
import math
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any, Literal
from fastapi import FastAPI, Request
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import JSONResponse
from pydantic import BaseModel, ConfigDict, ValidationError
@@ -212,10 +213,19 @@ def combine(d: Decision, results: list[dict]) -> dict:
"orderings": results}
def create_app(settings: Settings, engine: Any) -> FastAPI:
def inference_thread() -> ThreadPoolExecutor:
"""INV-2: the ONE host thread that ever touches the model. torch keeps CUDA state per host thread
(cuBLAS handles and workspaces), partly outside the VRAM cap: measured 2026-09-30, 40 worker
threads added 252 MiB outside the cap and 326 MiB inside it."""
return ThreadPoolExecutor(max_workers=1, thread_name_prefix="inference")
def create_app(settings: Settings, engine: Any, executor: ThreadPoolExecutor | None = None) -> FastAPI:
"""`executor` must be the single thread the engine was loaded on (main.py); tests get a fresh one."""
app = FastAPI(title="intern-decision-serve")
expected = f"Bearer {settings.api_token}".encode()
inference = threading.Lock() # INV-2: one request's calls at a time, off the event loop
executor = executor or inference_thread()
inference = threading.Lock() # INV-2: one request's calls at a time (the executor has one thread too)
in_progress = 0 # POSTs admitted and not yet answered
def admit():
@@ -288,9 +298,9 @@ def create_app(settings: Settings, engine: Any) -> FastAPI:
return out, timing
async def score(state: State, decisions: list[Decision], waves: list[list[Slot]]):
"""Run one request's calls in a worker thread; map their failures to contract codes."""
"""Run one request's calls on the inference thread; map their failures to contract codes."""
try:
return await run_in_threadpool(run, state, decisions, waves)
return await asyncio.get_running_loop().run_in_executor(executor, run, state, decisions, waves)
except ValueError as exc: # the model's own validation, token limit, ...
raise ApiError(422, "invalid_request", str(exc)) from exc
except OutOfMemory as exc:
@@ -5,7 +5,7 @@ import os
from fastapi import FastAPI
from .app import create_app
from .app import create_app, inference_thread
from .config import Settings
@@ -17,4 +17,7 @@ def app_from_env() -> FastAPI:
os.environ["TRANSFORMERS_OFFLINE"] = "1"
from .engine import TorchEngine # torch loads only inside load(), never in the unit tests
return create_app(settings, TorchEngine.load(settings))
# INV-2: load, warm-up and every later call run on this one thread (per-thread CUDA state).
executor = inference_thread()
engine = executor.submit(TorchEngine.load, settings).result()
return create_app(settings, engine, executor)