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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user