35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
import os
|
|
from ahserver.webapp import webapp
|
|
from ahserver.serverenv import ServerEnv
|
|
from ahserver.configuredServer import add_startup
|
|
from longtasks.longtasks import LongTasks, schedule_once
|
|
from appPublic.log import debug
|
|
|
|
class DemucsTasks(LongTasks):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.gpu_id = int(os.environ.get('DEMUCS_GPU_ID', '5'))
|
|
async def process_task(self, payload, workid=None):
|
|
import json
|
|
if isinstance(payload, str): payload = json.loads(payload)
|
|
task_type = payload.get('task_type', '')
|
|
if task_type == 'separate':
|
|
from workers.separate import run_separate
|
|
return await run_separate(self, payload)
|
|
raise ValueError(f'Unknown task_type: {task_type}')
|
|
|
|
async def on_app_built(app):
|
|
env = ServerEnv()
|
|
lt = env.longtasks
|
|
if lt:
|
|
schedule_once(0.1, lt.run)
|
|
debug(f'Demucs longtasks worker started, GPU: {lt.gpu_id}')
|
|
|
|
def init():
|
|
env = ServerEnv()
|
|
env.longtasks = DemucsTasks('redis://127.0.0.1:6379', 'demucs', worker_cnt=1, stuck_seconds=600, max_age_hours=24)
|
|
add_startup(on_app_built)
|
|
|
|
if __name__ == '__main__':
|
|
webapp(init)
|