136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
"""Spawned model processes, reused across tasks and replaceable after a crash."""
|
|
import logging
|
|
import multiprocessing
|
|
import queue
|
|
import time
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def analysis_process(connection, model_path, workers):
|
|
try:
|
|
from wanzheng import SafeYOLO, analyze_directory
|
|
model = SafeYOLO(model_path)
|
|
connection.send({'ready': True})
|
|
while True:
|
|
task = connection.recv()
|
|
if task is None:
|
|
return
|
|
result = analyze_directory(task['input_dir'], task['task_id'], model,
|
|
workers, task.get('api_url'))
|
|
connection.send(result)
|
|
except EOFError:
|
|
pass
|
|
except BaseException as exc:
|
|
logger.exception('常驻分析进程异常')
|
|
try:
|
|
connection.send({'error': str(exc) or type(exc).__name__})
|
|
except (OSError, EOFError):
|
|
pass
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
class ResidentPool:
|
|
def __init__(self, model_path, concurrency=1, workers=4, api_url=None,
|
|
startup_timeout=180, task_timeout=3600, target=analysis_process):
|
|
if min(concurrency, workers, startup_timeout, task_timeout) <= 0:
|
|
raise ValueError('并发、线程数和超时必须大于 0')
|
|
self.model_path = str(model_path)
|
|
self.concurrency = concurrency
|
|
self.workers = workers
|
|
self.api_url = api_url
|
|
self.startup_timeout = startup_timeout
|
|
self.task_timeout = task_timeout
|
|
self.target = target
|
|
self.context = multiprocessing.get_context('spawn')
|
|
self.slots = queue.Queue()
|
|
|
|
def receive(self, slot, timeout, stop_event=None):
|
|
connection, process = slot
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if stop_event is not None and stop_event.is_set():
|
|
raise RuntimeError('服务停止，任务中断，请重新提交')
|
|
if connection.poll(0.1):
|
|
try:
|
|
return connection.recv()
|
|
except (EOFError, OSError) as exc:
|
|
raise RuntimeError('常驻分析进程连接中断') from exc
|
|
if not process.is_alive():
|
|
raise RuntimeError(f'常驻分析进程退出：{process.exitcode}')
|
|
raise RuntimeError('常驻分析进程等待超时')
|
|
|
|
@staticmethod
|
|
def dispose(slot):
|
|
if slot is None:
|
|
return
|
|
connection, process = slot
|
|
connection.close()
|
|
if process.is_alive():
|
|
process.terminate()
|
|
process.join(timeout=5)
|
|
if process.is_alive():
|
|
process.kill()
|
|
process.join()
|
|
process.close()
|
|
|
|
def spawn(self, stop_event=None):
|
|
parent, child = self.context.Pipe()
|
|
process = self.context.Process(target=self.target,
|
|
args=(child, self.model_path, self.workers))
|
|
try:
|
|
process.start()
|
|
except BaseException:
|
|
parent.close()
|
|
child.close()
|
|
process.close()
|
|
raise
|
|
child.close()
|
|
slot = parent, process
|
|
try:
|
|
message = self.receive(slot, self.startup_timeout, stop_event)
|
|
if not message.get('ready'):
|
|
raise RuntimeError(message.get('error', '模型预热失败'))
|
|
return slot
|
|
except BaseException:
|
|
self.dispose(slot)
|
|
raise
|
|
|
|
def start(self):
|
|
try:
|
|
for _ in range(self.concurrency):
|
|
self.slots.put(self.spawn())
|
|
except BaseException:
|
|
self.stop()
|
|
raise
|
|
|
|
def stop(self):
|
|
# TaskQueue joins its worker threads before closing idle slots.
|
|
while True:
|
|
try:
|
|
self.dispose(self.slots.get_nowait())
|
|
except queue.Empty:
|
|
break
|
|
|
|
def __call__(self, task, stop_event):
|
|
slot = self.slots.get()
|
|
try:
|
|
if stop_event.is_set():
|
|
raise RuntimeError('服务停止，任务中断，请重新提交')
|
|
if slot is None or not slot[1].is_alive():
|
|
self.dispose(slot)
|
|
slot = None
|
|
slot = self.spawn(stop_event)
|
|
slot[0].send(dict(task, api_url=self.api_url))
|
|
result = self.receive(slot, self.task_timeout, stop_event)
|
|
except BaseException:
|
|
self.dispose(slot)
|
|
slot = None
|
|
raise
|
|
finally:
|
|
self.slots.put(slot)
|
|
if result.get('error'):
|
|
raise RuntimeError(result['error'])
|
|
return result
|