| """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: |
| |
| 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() |
|
|
