Zordon/server.py

import logging
import multiprocessing
import os
import socket
import threading
from pathlib import Path

import psutil

from src.api.api_server import API_VERSION
from src.api.api_server import start_api_server
from src.api.preview_manager import PreviewManager
from src.api.serverproxy_manager import ServerProxyManager
from src.application_context import ApplicationContext
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 (get_gpu_info, current_system_cpu, current_system_os,
                                       current_system_os_version, current_system_cpu_brand)
from src.utilities.zeroconf_server import ZeroconfServer
from src.version import APP_NAME, APP_VERSION

logger = logging.getLogger()


class ZordonServer:

    def __init__(self):
        self.ctx = ApplicationContext()

        # 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)

        # ---- Bootstrap Config ----
        Config.setup_config_dir()
        config_path = Path(Config.config_dir()) / "config.yaml"
        self.ctx.config = Config()
        self.ctx.config.load(config_path)
        Config._default_instance = self.ctx.config
        Config._sync_class()

        # ---- Engine Manager ----
        self.ctx.engine_manager = EngineManager()
        self.ctx.engine_manager.engines_path = Path(Config.upload_folder).expanduser() / "engines"
        os.makedirs(self.ctx.engine_manager.engines_path, exist_ok=True)
        EngineManager._default_instance = self.ctx.engine_manager
        EngineManager._sync_class()

        # ---- Preview Manager ----
        self.ctx.preview_manager = PreviewManager()
        self.ctx.preview_manager.storage_path = Path(Config.upload_folder).expanduser() / "previews"
        PreviewManager._default_instance = self.ctx.preview_manager
        PreviewManager._sync_class()

        # ---- Render Queue ----
        self.ctx.render_queue = RenderQueue()
        self.ctx.render_queue.load_state(database_directory=Path(Config.upload_folder).expanduser())
        RenderQueue._default_instance = self.ctx.render_queue
        RenderQueue._sync_class()

        # ---- Distributed Job Manager ----
        self.ctx.distributed_job_manager = DistributedJobManager()
        self.ctx.distributed_job_manager.subscribe_to_listener()
        DistributedJobManager._default_instance = self.ctx.distributed_job_manager
        DistributedJobManager._sync_class()

        self.api_server = None
        self.server_hostname: str = socket.gethostname()

    def start_server(self):

        def existing_process(process_name):
            current_pid = os.getpid()
            current_process = psutil.Process(current_pid)
            for proc in psutil.process_iter(['pid', 'name', 'ppid']):
                proc_name = proc.info['name'].lower().rstrip('.exe')
                if proc_name == process_name.lower() and proc.info['pid'] != current_pid:
                    if proc.info['pid'] == current_process.ppid():
                        continue  # parent process
                    elif proc.info['ppid'] == current_pid:
                        continue  # child process
                    else:
                        return proc  # unrelated process
            return None

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

        # main start
        logger.info(f"Starting {APP_NAME} Render Server ({APP_VERSION})")
        logger.debug(f"Upload directory: {Path(Config.upload_folder).expanduser()}")
        logger.debug(f"Thumbs directory: {PreviewManager.storage_path}")
        logger.debug(f"Engines directory: {EngineManager.engines_path}")

        ServerProxyManager.subscribe_to_listener()

        # update hostname
        self.server_hostname = socket.gethostname()

        # configure and start API server
        self.api_server = threading.Thread(target=start_api_server, args=(self.server_hostname,))
        self.api_server.daemon = True
        self.api_server.start()

        # start zeroconf server
        ctx = self.ctx
        ctx.zeroconf_server = ZeroconfServer()
        ctx.zeroconf_server._configure(f"_{APP_NAME.lower()}._tcp.local.", self.server_hostname, Config.port_number)
        ctx.zeroconf_server.properties = {'system_cpu': current_system_cpu(),
                                          'system_cpu_brand': current_system_cpu_brand(),
                                          'system_cpu_cores': multiprocessing.cpu_count(),
                                          'system_os': current_system_os(),
                                          'system_os_version': current_system_os_version(),
                                          'system_memory': round(psutil.virtual_memory().total / (1024**3)),
                                          'gpu_info': get_gpu_info(),
                                          'api_version': API_VERSION}
        ZeroconfServer._default_instance = ctx.zeroconf_server
        ZeroconfServer._sync_class()
        ctx.zeroconf_server._start()

        logger.info(f"{APP_NAME} Render Server started - Hostname: {self.server_hostname}")
        RenderQueue.start()  # Start evaluating the render queue

    def is_running(self):
        return self.api_server and self.api_server.is_alive()

    def stop_server(self):
        logger.info(f"{APP_NAME} Render Server is preparing to stop")
        try:
            ZeroconfServer.stop()
            RenderQueue.prepare_for_shutdown()
        except Exception as e:
            logger.exception(f"Exception during prepare for shutdown: {e}")
        logger.info(f"{APP_NAME} Render Server has shut down")


if __name__ == '__main__':
    server = ZordonServer()
    try:
        server.start_server()
        server.api_server.join()
    except KeyboardInterrupt:
        pass
    except Exception as e:
        logger.fatal(f"Unhandled exception: {e}")
    finally:
        server.stop_server()