trap/trap/plumber.py
2024-12-03 15:21:52 +01:00

148 lines
No EOL
5.1 KiB
Python

import atexit
import logging
from logging.handlers import SocketHandler, QueueHandler, QueueListener
from multiprocessing import Event, Process, Queue
import multiprocessing
import signal
import sys
import time
from trap.config import parser
from trap.cv_renderer import run_cv_renderer
from trap.frame_emitter import run_frame_emitter
from trap.prediction_server import run_prediction_server
from trap.preview_renderer import run_preview_renderer
from trap.animation_renderer import run_animation_renderer
from trap.socket_forwarder import run_ws_forwarder
from trap.tracker import run_tracker
from setproctitle import setproctitle, setthreadtitle
logger = logging.getLogger("trap.plumbing")
class ExceptionHandlingProcess(Process):
def run(self):
assert 'is_running' in self._kwargs
# exit handler to make sure that on many kinds of crashes and kills the
# suprocess warns parent/siblings
# TODO: Does not work with OOM kill. Would need a watchdog process for that
def exit_handler(*args):
self._kwargs['is_running'].clear()
atexit.register(exit_handler)
signal.signal(signal.SIGTERM, exit_handler)
signal.signal(signal.SIGINT, exit_handler)
setproctitle(f"trap-{self.name}")
try:
super(Process, self).run()
print("finished ", self.name)
except BaseException as e:
logger.critical(f"Exception in {self.name}")
logger.exception(e)
self._kwargs['is_running'].clear()
def start():
args = parser.parse_args()
loglevel = logging.NOTSET if args.verbose > 1 else logging.DEBUG if args.verbose > 0 else logging.INFO
# print(args)
# exit()
isRunning = Event()
isRunning.set()
q = multiprocessing.Queue(-1)
queue_handler = QueueHandler(q)
stream_handler = logging.StreamHandler()
log_handlers = [stream_handler]
if args.remote_log_addr:
logging.captureWarnings(True)
# root_logger.setLevel(logging.NOTSET) # to send all records to cutelog
socket_handler = SocketHandler(args.remote_log_addr, args.remote_log_port)
socket_handler.setLevel(logging.NOTSET)
log_handlers.append(socket_handler)
queue_listener = QueueListener(q, *log_handlers, respect_handler_level=True)
# root = logging.getLogger()
logging.basicConfig(
level=loglevel,
handlers=[queue_handler]
)
# root_logger = logging.getLogger()
# # set per handler, so we can set it lower for the root logger if remote logging is enabled
# [h.setLevel(loglevel) for h in root_logger.handlers]
# queue_listener.handlers.append(socket_handler)
# instantiating process with arguments
procs = [
# ExceptionHandlingProcess(target=run_ws_forwarder, kwargs={'config': args, 'is_running': isRunning}, name='forwarder'),
ExceptionHandlingProcess(target=run_frame_emitter, kwargs={'config': args, 'is_running': isRunning}, name='frame_emitter'),
ExceptionHandlingProcess(target=run_tracker, kwargs={'config': args, 'is_running': isRunning}, name='tracker'),
]
if args.render_file or args.render_url or args.render_window:
if not args.render_no_preview or args.render_file or args.render_url:
procs.append(
# ExceptionHandlingProcess(target=run_cv_renderer, kwargs={'config': args, 'is_running': isRunning}, name='preview')
ExceptionHandlingProcess(target=run_cv_renderer, kwargs={'config': args, 'is_running': isRunning}, name='preview')
)
if args.render_no_preview:
procs.append(
ExceptionHandlingProcess(target=run_animation_renderer, kwargs={'config': args, 'is_running': isRunning}, name='renderer')
)
if not args.bypass_prediction:
procs.append(
ExceptionHandlingProcess(target=run_prediction_server, kwargs={'config': args, 'is_running':isRunning}, name='inference'),
)
try:
logger.info("start")
for proc in procs:
proc.start()
# if start the listener before the subprocesses, it becomes a mess, because the
# running threat is forked too, but cannot easily be stopped in the forks.
# Thus, only start the queue-listener threat _after_ starting processes
queue_listener.start()
# wait for processes to clean up
for proc in procs:
proc.join()
isRunning.clear()
logger.info('Stop')
except BaseException as e:
# mainly for KeyboardInterrupt
# but in any case, on error all processed need to be signalled to shut down
logger.critical(f"Exception in plumber")
logger.exception(e)
isRunning.clear()
# while True:
# time.sleep(2)
# any_alive = False
# alive = [proc for proc in procs if proc.is_alive()]
# print("alive: ", [p.name for p in alive])
# if len(alive) < 1:
# break
print('stop listener')
queue_listener.stop()
print('stopped listener')
print("finished plumber")
if __name__ == "__main__":
start()