Zordon/src/init.py

import logging
import multiprocessing
import os
import socket
import sys
import threading
from collections import deque

from src.api.api_server import start_server
from src.api.preview_manager import PreviewManager
from src.api.serverproxy_manager import ServerProxyManager
from src.distributed_job_manager import DistributedJobManager
from src.engines.engine_manager import EngineManager
from src.render_queue import RenderQueue
from src.utilities.config import Config
from src.utilities.misc_helper import system_safe_path, current_system_cpu, current_system_os, current_system_os_version
from src.utilities.zeroconf_server import ZeroconfServer
from version import APP_NAME

logger = logging.getLogger()


def run(server_only=False) -> int:
    """Initializes the application and runs it.

    Args:
        server_only: Run in server-only CLI mode. Default is False (runs in GUI mode).

    Returns:
        int: The exit status code.
    """

    def existing_process(process_name):
        import psutil
        for proc in psutil.process_iter(['pid', 'name']):
            if proc.info['name'].lower() == process_name.lower() and proc.info['pid'] != os.getpid():
                return proc
        return None

    # setup logging
    logging.basicConfig(format='%(asctime)s: %(levelname)s: %(module)s: %(message)s', datefmt='%d-%b-%y %H:%M:%S',
                        level=Config.server_log_level.upper())
    logging.getLogger("requests").setLevel(logging.WARNING)  # suppress noisy requests/urllib3 logging
    logging.getLogger("urllib3").setLevel(logging.WARNING)

    # check for existing instance
    existing_proc = existing_process(APP_NAME)
    if existing_proc:
        logger.fatal(f"Another instance of {APP_NAME} is already running (pid: {existing_proc.pid})")
        sys.exit(1)

    # Setup logging for console ui
    buffer_handler = __setup_buffer_handler() if not server_only else None

    logger.info(f"Starting {APP_NAME} Render Server")
    return_code = 0
    try:
        # Load Config YAML
        Config.setup_config_dir()
        Config.load_config(system_safe_path(os.path.join(Config.config_dir(), 'config.yaml')))

        # configure default paths
        EngineManager.engines_path = system_safe_path(
            os.path.join(os.path.join(os.path.expanduser(Config.upload_folder),
                                      'engines')))
        os.makedirs(EngineManager.engines_path, exist_ok=True)
        PreviewManager.storage_path = system_safe_path(
            os.path.join(os.path.expanduser(Config.upload_folder), 'previews'))

        # Debug info
        logger.debug(f"Upload directory: {os.path.expanduser(Config.upload_folder)}")
        logger.debug(f"Thumbs directory: {PreviewManager.storage_path}")
        logger.debug(f"Engines directory: {EngineManager.engines_path}")

        # Set up the RenderQueue object
        RenderQueue.load_state(database_directory=system_safe_path(os.path.expanduser(Config.upload_folder)))
        ServerProxyManager.subscribe_to_listener()
        DistributedJobManager.subscribe_to_listener()

        # check for updates for render engines if configured or on first launch
        if Config.update_engines_on_launch or not EngineManager.get_engines():
            EngineManager.update_all_engines()

        # get hostname
        local_hostname = socket.gethostname()
        local_hostname = local_hostname + (".local" if not local_hostname.endswith(".local") else "")

        # configure and start API server
        api_server = threading.Thread(target=start_server, args=(local_hostname,))
        api_server.daemon = True
        api_server.start()

        # start zeroconf server
        ZeroconfServer.configure(f"_{APP_NAME.lower()}._tcp.local.", local_hostname, Config.port_number)
        ZeroconfServer.properties = {'system_cpu': current_system_cpu(),
                                     'system_cpu_cores': multiprocessing.cpu_count(),
                                     'system_os': current_system_os(),
                                     'system_os_version': current_system_os_version()}
        ZeroconfServer.start()
        logger.info(f"{APP_NAME} Render Server started - Hostname: {local_hostname}")
        RenderQueue.start()  # Start evaluating the render queue

        # start in gui or server only (cli) mode
        logger.debug(f"Launching in {'server only' if server_only else 'GUI'} mode")
        if server_only:  # CLI only
            api_server.join()
        else:  # GUI
            return_code = __show_gui(buffer_handler)

    except KeyboardInterrupt:
        pass
    except Exception as e:
        logging.error(f"Unhandled exception: {e}")
        return_code = 1
    finally:
        # shut down gracefully
        logger.info(f"{APP_NAME} Render Server is preparing to shut down")
        try:
            RenderQueue.prepare_for_shutdown()
        except Exception as e:
            logger.exception(f"Exception during prepare for shutdown: {e}")
        ZeroconfServer.stop()
        logger.info(f"{APP_NAME} Render Server has shut down")
        return sys.exit(return_code)


def __setup_buffer_handler():
    # lazy load GUI frameworks
    from PyQt6.QtCore import QObject, pyqtSignal

    class BufferingHandler(logging.Handler, QObject):
        new_record = pyqtSignal(str)

        def __init__(self, capacity=100):
            logging.Handler.__init__(self)
            QObject.__init__(self)
            self.buffer = deque(maxlen=capacity)  # Define a buffer with a fixed capacity

        def emit(self, record):
            try:
                msg = self.format(record)
                self.buffer.append(msg)  # Add message to the buffer
                self.new_record.emit(msg)  # Emit signal
            except RuntimeError:
                pass

        def get_buffer(self):
            return list(self.buffer)  # Return a copy of the buffer

    buffer_handler = BufferingHandler()
    buffer_handler.setFormatter(logging.getLogger().handlers[0].formatter)
    new_logger = logging.getLogger()
    new_logger.addHandler(buffer_handler)
    return buffer_handler


def __show_gui(buffer_handler):
    # lazy load GUI frameworks
    from PyQt6.QtWidgets import QApplication

    # load application
    app: QApplication = QApplication(sys.argv)

    # configure main window
    from src.ui.main_window import MainWindow
    window: MainWindow = MainWindow()
    window.buffer_handler = buffer_handler
    window.show()

    return app.exec()