Spaces:
Sleeping
Sleeping
| from speakers.common.registry import registry | |
| from speakers.server.bootstrap.bootstrap_register import load_bootstrap, get_bootstrap | |
| from omegaconf import OmegaConf | |
| from speakers.common.utils import get_abs_path | |
| from oscrypto import util as crypto_utils | |
| import asyncio | |
| import time | |
| import os | |
| import sys | |
| import traceback | |
| import subprocess | |
| root_dir = os.path.dirname(os.path.abspath(__file__)) | |
| registry.register_path("server_library_root", root_dir) | |
| # Time to wait for web client to send a request to /task-state request | |
| # before that web clients task gets removed from the queue | |
| WEB_CLIENT_TIMEOUT = -1 | |
| # Time before finished tasks get removed from memory | |
| FINISHED_TASK_REMOVE_TIMEOUT = 1800 | |
| def generate_nonce(): | |
| return crypto_utils.rand_bytes(16).hex() | |
| def start_translator_client_proc(speakers_config_file: str, nonce: str = None): | |
| cmds = [ | |
| sys.executable, | |
| '-m', 'speakers', | |
| '--mode', 'web_runner', | |
| '--speakers-config-file', speakers_config_file, | |
| '--nonce', nonce, | |
| ] | |
| proc = subprocess.Popen(cmds, cwd=f"{registry.get_path('library_root')}/../") | |
| return proc | |
| async def start_async_app(speakers_config_file: str, nonce: str = None): | |
| config = OmegaConf.load(get_abs_path(speakers_config_file)) | |
| load_bootstrap(config=config.get("bootstrap")) | |
| runner_bootstrap_web = get_bootstrap("runner_bootstrap_web") | |
| runner_bootstrap_web.set_nonce(nonce=nonce) | |
| await runner_bootstrap_web.run() | |
| return runner_bootstrap_web | |
| async def dispatch(speakers_config_file: str, nonce: str = None): | |
| if nonce is None: | |
| nonce = os.getenv('MT_WEB_NONCE', generate_nonce()) | |
| runner = await start_async_app(speakers_config_file=speakers_config_file, nonce=nonce) | |
| # Create client process | |
| print() | |
| client_process = start_translator_client_proc(speakers_config_file, nonce=nonce) | |
| try: | |
| while True: | |
| """任务队列状态维护""" | |
| await asyncio.sleep(1) | |
| # Restart client if OOM or similar errors occured | |
| if client_process.poll() is not None: | |
| print('Restarting translator process') | |
| if len(runner.ongoing_tasks) > 0: | |
| task_id = runner.ongoing_tasks.pop(0) | |
| state = runner.task_states[task_id] | |
| state['info'] = 'error' | |
| state['finished'] = True | |
| client_process = start_translator_client_proc(speakers_config_file=speakers_config_file) | |
| # Filter queued and finished tasks | |
| now = time.time() | |
| to_del_task_ids = set() | |
| for tid, s in runner.task_states.items(): | |
| payload = runner.task_data[tid] | |
| # Remove finished tasks after 30 minutes | |
| if s['finished'] and now - payload.created_at > FINISHED_TASK_REMOVE_TIMEOUT: | |
| to_del_task_ids.add(tid) | |
| # Remove queued tasks without web client | |
| elif WEB_CLIENT_TIMEOUT >= 0: | |
| if tid not in runner.ongoing_tasks and not s['finished'] \ | |
| and now - payload.requested_at > WEB_CLIENT_TIMEOUT: | |
| print('REMOVING TASK', tid) | |
| to_del_task_ids.add(tid) | |
| try: | |
| runner.queue.remove(tid) | |
| except Exception: | |
| pass | |
| for tid in to_del_task_ids: | |
| del runner.task_states[tid] | |
| del runner.task_data[tid] | |
| except: | |
| if client_process.poll() is None: | |
| # client_process.terminate() | |
| client_process.kill() | |
| await runner.destroy() | |
| traceback.print_exc() | |
| raise | |