63 lines
2.7 KiB
Python
63 lines
2.7 KiB
Python
"""CMP opt-in: one continuous chapter worker per GPU, with immediate progress."""
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import httpx
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_lock = asyncio.Lock()
|
|
|
|
|
|
async def warm_chapters(script, count, endpoints, params, disconnected=None):
|
|
if count <= 0:
|
|
return
|
|
if not endpoints:
|
|
raise ValueError('No CMP GPU endpoints configured')
|
|
async with _lock:
|
|
if disconnected and await disconnected():
|
|
raise asyncio.CancelledError()
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(7200, connect=15), trust_env=False) as client:
|
|
indices = iter(range(count))
|
|
results = asyncio.Queue()
|
|
|
|
async def worker(endpoint):
|
|
try:
|
|
for index in indices:
|
|
payload = dict(params, text=script, chapter_index=index)
|
|
response = await client.post(endpoint + '/audiobook/preview', json=payload)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
logger.info('CMP_GPU_CHAPTER gpu=%s chapter=%s cached=%s duration=%s',
|
|
endpoint, index, result.get('cached'), result.get('duration_s'))
|
|
event = dict(type='chapter', index=index, total=count,
|
|
title=result.get('title', ''), duration_s=result['duration_s'],
|
|
cached=result.get('cached', False), stage='gpu_render')
|
|
await results.put(('chapter', event))
|
|
except Exception as error:
|
|
await results.put(('error', error))
|
|
finally:
|
|
await results.put(('finished', None))
|
|
|
|
tasks = [asyncio.create_task(worker(endpoint)) for endpoint in endpoints[:count]]
|
|
finished = 0
|
|
try:
|
|
while finished < len(tasks):
|
|
if disconnected and await disconnected():
|
|
raise asyncio.CancelledError()
|
|
try:
|
|
kind, value = await asyncio.wait_for(results.get(), timeout=5)
|
|
except asyncio.TimeoutError:
|
|
yield ': cmp GPU rendering\n\n'
|
|
continue
|
|
if kind == 'error':
|
|
raise value
|
|
if kind == 'finished':
|
|
finished += 1
|
|
else:
|
|
yield 'data: ' + json.dumps(value) + '\n\n'
|
|
finally:
|
|
for task in tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|