File size: 1,821 Bytes
5d6f353 | 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 | """Observable, single-flight model warmup without importing CUDA on HTTP import."""
import asyncio
import threading
import time
class ModelRuntime:
def __init__(self, loader=None):
self.loader = loader
self._lock = threading.Lock()
self._task = None
self._state = {'state': 'not_started', 'model_ready': False,
'load_seconds': None, 'error': None, 'devices': []}
def snapshot(self):
with self._lock:
return {**self._state, 'devices': [dict(d) for d in self._state['devices']]}
def _load(self):
started = time.monotonic()
try:
if self.loader:
devices = self.loader() or []
else:
from llm import _load, gpu_report
_load()
devices = gpu_report()
result = {'state': 'ready', 'model_ready': True, 'error': None, 'devices': devices}
except Exception as exc:
# Exception messages can contain cache paths or credentials in URLs.
result = {'state': 'failed', 'model_ready': False,
'error': type(exc).__name__, 'devices': []}
print('[LLM] Warmup failed: ' + type(exc).__name__, flush=True)
with self._lock:
self._state.update(result, load_seconds=time.monotonic() - started)
def start(self):
if self._task is None:
with self._lock:
self._state.update(state='loading', model_ready=False)
self._task = asyncio.create_task(asyncio.to_thread(self._load))
return self._task
async def wait_ready(self):
await asyncio.shield(self.start())
if not self.snapshot()['model_ready']:
raise RuntimeError('model_unavailable')
MODEL_RUNTIME = ModelRuntime()
|