voicebox/backend/server.py

"""
Entry point for PyInstaller-bundled voicebox server.

This module provides an entry point that works with PyInstaller by using
absolute imports instead of relative imports.
"""

import sys

# Fast path: handle --version before any heavy imports so the Rust
# version check doesn't block for 30+ seconds loading torch etc.
if "--version" in sys.argv:
    from backend import __version__
    print(f"voicebox-server {__version__}")
    sys.exit(0)

import logging

# Set up logging FIRST, before any imports that might fail
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    stream=sys.stderr,  # Log to stderr so it's captured by Tauri
)
logger = logging.getLogger(__name__)

# Log startup immediately to confirm binary execution
logger.info("=" * 60)
logger.info("voicebox-server starting up...")
logger.info(f"Python version: {sys.version}")
logger.info(f"Executable: {sys.executable}")
logger.info(f"Arguments: {sys.argv}")
logger.info("=" * 60)

try:
    logger.info("Importing argparse...")
    import argparse
    logger.info("Importing uvicorn...")
    import uvicorn
    logger.info("Standard library imports successful")

    # Import the FastAPI app from the backend package
    logger.info("Importing backend.config...")
    from backend import config
    logger.info("Importing backend.database...")
    from backend import database
    logger.info("Importing backend.main (this may take a while due to torch/transformers)...")
    from backend.main import app
    logger.info("Backend imports successful")
except Exception as e:
    logger.error(f"Failed to import required modules: {e}", exc_info=True)
    sys.exit(1)

def _start_parent_watchdog(parent_pid):
    """Monitor parent process and exit if it dies.

    This is the clean shutdown mechanism: instead of the Tauri app trying to
    forcefully kill the server (which spawns console windows on Windows),
    the server monitors its parent and shuts itself down gracefully.
    """
    import os
    import signal
    import threading
    import time

    def _is_pid_alive(pid):
        """Check if a process with the given PID exists (cross-platform)."""
        try:
            if sys.platform == "win32":
                import ctypes
                kernel32 = ctypes.windll.kernel32
                SYNCHRONIZE = 0x00100000
                handle = kernel32.OpenProcess(SYNCHRONIZE, False, pid)
                if handle:
                    kernel32.CloseHandle(handle)
                    return True
                return False
            else:
                os.kill(pid, 0)
                return True
        except (OSError, PermissionError):
            return False

    def _watch():
        logger.info(f"Parent watchdog started, monitoring PID {parent_pid}")
        while True:
            if not _is_pid_alive(parent_pid):
                logger.info(f"Parent process {parent_pid} no longer exists, shutting down...")
                os.kill(os.getpid(), signal.SIGTERM)
                return
            time.sleep(1)

    t = threading.Thread(target=_watch, daemon=True)
    t.start()


if __name__ == "__main__":
    try:
        parser = argparse.ArgumentParser(description="voicebox backend server")
        parser.add_argument(
            "--host",
            type=str,
            default="127.0.0.1",
            help="Host to bind to (use 0.0.0.0 for remote access)",
        )
        parser.add_argument(
            "--port",
            type=int,
            default=8000,
            help="Port to bind to",
        )
        parser.add_argument(
            "--data-dir",
            type=str,
            default=None,
            help="Data directory for database, profiles, and generated audio",
        )
        parser.add_argument(
            "--parent-pid",
            type=int,
            default=None,
            help="PID of parent process to monitor; server exits when parent dies",
        )
        parser.add_argument(
            "--version",
            action="store_true",
            help="Print version and exit (handled above, kept for argparse help)",
        )
        args = parser.parse_args()

        # Detect backend variant from binary name
        # voicebox-server-cuda → sets VOICEBOX_BACKEND_VARIANT=cuda
        import os
        binary_name = os.path.basename(sys.executable).lower()
        if "cuda" in binary_name:
            os.environ["VOICEBOX_BACKEND_VARIANT"] = "cuda"
            logger.info("Backend variant: CUDA")
        else:
            os.environ["VOICEBOX_BACKEND_VARIANT"] = "cpu"
            logger.info("Backend variant: CPU")

        # Start parent process watchdog if requested
        if args.parent_pid is not None:
            _start_parent_watchdog(args.parent_pid)

        logger.info(f"Parsed arguments: host={args.host}, port={args.port}, data_dir={args.data_dir}")

        # Set data directory if provided
        if args.data_dir:
            logger.info(f"Setting data directory to: {args.data_dir}")
            config.set_data_dir(args.data_dir)

        # Initialize database after data directory is set
        logger.info("Initializing database...")
        database.init_db()
        logger.info("Database initialized successfully")

        logger.info(f"Starting uvicorn server on {args.host}:{args.port}...")
        uvicorn.run(
            app,
            host=args.host,
            port=args.port,
            log_level="info",
        )
    except Exception as e:
        logger.error(f"Server startup failed: {e}", exc_info=True)
        sys.exit(1)