Zordon/server.py

import logging
import multiprocessing
import os
import socket
import sys
import threading

import cpuinfo
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.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, system_safe_path, current_system_cpu, current_system_os,
                                       current_system_os_version, current_system_cpu_brand, check_for_updates)
from src.utilities.zeroconf_server import ZeroconfServer
from src.version import APP_NAME, APP_VERSION, APP_REPO_NAME, APP_REPO_OWNER

logger = logging.getLogger()


def start_server(skip_updates=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
        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

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

    # check for updates
    if not skip_updates:
        update_thread = threading.Thread(target=check_for_updates, args=(APP_REPO_NAME, APP_REPO_OWNER, APP_NAME,
                                                                         APP_VERSION))
        update_thread.start()

    # main start
    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()

        # get hostname
        local_hostname = socket.gethostname()

        # configure and start API server
        api_server = threading.Thread(target=start_api_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_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)),  # in GB
                                     'gpu_info': get_gpu_info(),
                                     'api_version': API_VERSION}
        ZeroconfServer.start()
        logger.info(f"{APP_NAME} Render Server started - Hostname: {local_hostname}")
        RenderQueue.start()  # Start evaluating the render queue

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

        api_server.join()

    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)

if __name__ == '__main__':
    start_server()