faceswap/lib/queue_manager.py

#!/usr/bin/env python3
""" Queue Manager for faceswap

    NB: Keep this in it's own module! If it gets loaded from
    a multiprocess on a Windows System it will break Faceswap"""

import logging
import multiprocessing as mp
import sys
import threading

from queue import Queue, Empty as QueueEmpty  # pylint: disable=unused-import; # noqa
from time import sleep

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name


class QueueManager():
    """ Manage queues for availabilty across processes
        Don't import this class directly, instead
        import the variable: queue_manager """
    def __init__(self):
        logger.debug("Initializing %s", self.__class__.__name__)

        # Hacky fix to stop multiprocessing spawning managers in child processes
        if mp.current_process().name == "MainProcess":
            # Use a Multiprocessing manager in main process
            self.manager = mp.Manager()
        else:
            # Use a standard mp.queue in child process. NB: This will never be used
            # but spawned processes will load this module, so we need to dummy in a queue
            self.manager = mp
        self.shutdown = self.manager.Event()
        self.queues = dict()
        # Despite launching a subprocess, the scripts still want to access the same logging
        # queue as the GUI, so make sure the GUI gets it's own queue
        self._log_queue = self.manager.Queue() if "gui" not in sys.argv else mp.Queue()
        logger.debug("Initialized %s", self.__class__.__name__)

    def add_queue(self, name, maxsize=0, multiprocessing_queue=True):
        """ Add a queue to the manager

            Adds an event "shutdown" to the queue that can be used to indicate
            to a process that any activity on the queue should cease """

        logger.debug("QueueManager adding: (name: '%s', maxsize: %s)", name, maxsize)
        if name in self.queues.keys():
            raise ValueError("Queue '{}' already exists.".format(name))

        if multiprocessing_queue:
            queue = self.manager.Queue(maxsize=maxsize)
        else:
            queue = Queue(maxsize=maxsize)

        setattr(queue, "shutdown", self.shutdown)
        self.queues[name] = queue
        logger.debug("QueueManager added: (name: '%s')", name)

    def del_queue(self, name):
        """ remove a queue from the manager """
        logger.debug("QueueManager deleting: '%s'", name)
        del self.queues[name]
        logger.debug("QueueManager deleted: '%s'", name)

    def get_queue(self, name, maxsize=0):
        """ Return a queue from the manager
            If it doesn't exist, create it """
        logger.debug("QueueManager getting: '%s'", name)
        queue = self.queues.get(name, None)
        if not queue:
            self.add_queue(name, maxsize)
            queue = self.queues[name]
        logger.debug("QueueManager got: '%s'", name)
        return queue

    def terminate_queues(self):
        """ Set shutdown event, clear and send EOF to all queues
            To be called if there is an error """
        logger.debug("QueueManager terminating all queues")
        self.shutdown.set()
        for q_name, queue in self.queues.items():
            logger.debug("QueueManager terminating: '%s'", q_name)
            while not queue.empty():
                queue.get(True, 1)
            queue.put("EOF")
        logger.debug("QueueManager terminated all queues")

    def debug_monitor(self, update_secs=2):
        """ Debug tool for monitoring queues """
        thread = threading.Thread(target=self.debug_queue_sizes,
                                  args=(update_secs, ))
        thread.daemon = True
        thread.start()

    def debug_queue_sizes(self, update_secs):
        """ Output the queue sizes
            logged to INFO so it also displays in console
        """
        while True:
            for name in sorted(self.queues.keys()):
                logger.info("%s: %s", name, self.queues[name].qsize())
            sleep(update_secs)


queue_manager = QueueManager()  # pylint: disable=invalid-name