mirror of
https://github.com/iperov/DeepFaceLive.git
synced 2024-12-25 07:21:13 -08:00
143 lines
4.2 KiB
Python
143 lines
4.2 KiB
Python
import multiprocessing
|
|
import threading
|
|
import time
|
|
import traceback
|
|
import weakref
|
|
from typing import List
|
|
|
|
class MPWorker:
|
|
def __init__(self, sub_args : List = None,
|
|
process_count : int = None ):
|
|
"""
|
|
base class for multi process worker
|
|
|
|
provides messaging interface between host and subprocesses
|
|
|
|
|
|
sub_args a list of args will be passed to _on_sub_initialize
|
|
|
|
process_count number of subprocesses. Default : number of cpu count
|
|
|
|
starts immediatelly after construction.
|
|
"""
|
|
if process_count is None:
|
|
process_count = multiprocessing.cpu_count()
|
|
|
|
pipes = []
|
|
ps = []
|
|
for i in range(process_count):
|
|
host_pipe, sub_pipe = multiprocessing.Pipe()
|
|
|
|
p = multiprocessing.Process(target=self._sub_process, args=(i, process_count, sub_pipe, sub_args), daemon=True)
|
|
p.start()
|
|
|
|
pipes += [host_pipe]
|
|
ps += [p]
|
|
|
|
self._process_id = -1
|
|
self._process_count = process_count
|
|
self._process_working_count = process_count
|
|
self._pipes = pipes
|
|
self._ps = ps
|
|
|
|
threading.Thread(target=_host_thread_proc, args=(weakref.ref(self),), daemon=True).start()
|
|
|
|
# overridable
|
|
def _on_host_sub_message(self, process_id, name, *args, **kwargs):
|
|
"""a message from subprocess"""
|
|
# overridable
|
|
def _on_sub_host_message(self, name, *args, **kwargs):
|
|
"""a message from host"""
|
|
# overridable
|
|
def _on_sub_initialize(self, *args):
|
|
"""on subprocess initialization"""
|
|
# overridable
|
|
def _on_sub_finalize(self):
|
|
"""on graceful subprocess finalization"""
|
|
# overridable
|
|
def _on_sub_tick(self, process_id):
|
|
""""""
|
|
def get_process_count(self) -> int: return self._process_count
|
|
def get_process_id(self) -> int: return self._process_id
|
|
|
|
def kill(self):
|
|
"""
|
|
kill subprocess
|
|
"""
|
|
for p in self._ps:
|
|
p.kill()
|
|
self._ps = []
|
|
|
|
def stop(self):
|
|
"""
|
|
graceful stop subprocess, will wait all subprocess finalization
|
|
"""
|
|
self._send_msg('__stop')
|
|
for p in self._ps:
|
|
p.join()
|
|
self._ps = []
|
|
|
|
def _host_process_messages(self, timeout : float = 0) -> bool:
|
|
"""
|
|
process messages on host side
|
|
"""
|
|
for process_id, pipe in enumerate(self._pipes):
|
|
try:
|
|
if pipe.poll(timeout):
|
|
name, args, kwargs = pipe.recv()
|
|
if name == '__stopped':
|
|
self._process_working_count -= 1
|
|
else:
|
|
self._on_host_sub_message(process_id, name, *args, **kwargs)
|
|
except:
|
|
...
|
|
|
|
def _send_msg(self, name, *args, process_id=-1, **kwargs):
|
|
"""
|
|
send message to other side
|
|
|
|
process_id -1 mean send to all sub processes
|
|
on subprocess side - ignore this param
|
|
"""
|
|
try:
|
|
for i, pipe in enumerate(self._pipes):
|
|
if process_id == -1 or i == process_id:
|
|
pipe.send( (name, args, kwargs) )
|
|
except:
|
|
...
|
|
|
|
def _sub_process(self, process_id, process_count, pipe, sub_args):
|
|
self._process_id = process_id
|
|
self._process_count = process_count
|
|
self._pipes = [pipe]
|
|
|
|
self._on_sub_initialize(*sub_args)
|
|
|
|
working = True
|
|
while working:
|
|
self._on_sub_tick(process_id)
|
|
|
|
if pipe.poll(0.005):
|
|
while True:
|
|
name, args, kwargs = pipe.recv()
|
|
if name == '__stop':
|
|
working = False
|
|
else:
|
|
try:
|
|
self._on_sub_host_message(name, *args, **kwargs)
|
|
except:
|
|
print(f'Error during handling host message {name} : {traceback.format_exc()}')
|
|
|
|
if not pipe.poll():
|
|
break
|
|
|
|
self._on_sub_finalize()
|
|
|
|
def _host_thread_proc(wref):
|
|
while True:
|
|
ref = wref()
|
|
if ref is None:
|
|
break
|
|
ref._host_process_messages(0)
|
|
del ref
|
|
time.sleep(0.005) |