DeblurGanV2/run.py

import os
import subprocess
import sys
import time
import webbrowser
import requests
import signal
import psutil

def main():
    """Run the Debluring Application by starting both API and Streamlit app"""
    # Change to the directory of this script
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    
    print("\n🚀 Starting Debluring Application...\n")
    
    # Define URLs
    api_url = "http://localhost:8001"
    streamlit_url = "http://localhost:8501"
    
    # Function to start or restart the API server
    def start_api_server():
        nonlocal api_process
        # Kill any existing uvicorn processes that might be hanging
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                if 'uvicorn' in str(proc.info['cmdline']).lower() and '8001' in str(proc.info['cmdline']):
                    print(f"Killing existing uvicorn process (PID: {proc.info['pid']})")
                    psutil.Process(proc.info['pid']).kill()
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
                
        print("Starting API server...")
        return subprocess.Popen(
            [sys.executable, "api.py"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            bufsize=1,
            universal_newlines=True
        )
    
    # Function to check API health
    def check_api_health(timeout=2):
        try:
            response = requests.get(f"{api_url}/status", timeout=timeout)
            return response.status_code == 200
        except:
            return False
    
    # Start the API server
    api_process = start_api_server()
    
    # Wait for API to actually be available
    print("Waiting for API to start", end="")
    api_ready = False
    for _ in range(15):  # Try for 15 seconds
        if check_api_health():
            api_ready = True
            print("\n✅ API server is running")
            break
        print(".", end="", flush=True)
        time.sleep(1)
        
    if not api_ready:
        print("\n⚠️ API server might not be fully ready, but continuing anyway")
    
    # Function to start or restart the Streamlit app
    def start_streamlit_app():
        nonlocal streamlit_process
        # Kill any existing streamlit processes that might be hanging
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                if 'streamlit' in str(proc.info['cmdline']).lower() and 'run' in str(proc.info['cmdline']):
                    print(f"Killing existing streamlit process (PID: {proc.info['pid']})")
                    psutil.Process(proc.info['pid']).kill()
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
                
        print("Starting Streamlit web interface...")
        return subprocess.Popen(
            [sys.executable, "-m", "streamlit", "run", "app.py"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            bufsize=1,
            universal_newlines=True
        )
    
    # Start the Streamlit app
    streamlit_process = start_streamlit_app()
    
    # Give Streamlit a moment to start
    time.sleep(3)
    print("✅ Streamlit web interface started")
    
    # Open web browser
    try:
        print("Opening web interface in your browser...")
        webbrowser.open(streamlit_url)
    except Exception as e:
        print(f"Failed to open browser. Please open this URL manually: {streamlit_url}")
    
    # Print URLs
    print("\n📋 Application URLs:")
    print(f"   - Web Interface: {streamlit_url}")
    print(f"   - API: {api_url}")
    
    # Set up a more graceful exit handler
    def handle_exit(signum, frame):
        print("\n👋 Shutting down gracefully...")
        shutdown_services(api_process, streamlit_process)
        sys.exit(0)
        
    signal.signal(signal.SIGINT, handle_exit)
    
    print("\n⌨️  Press Ctrl+C to stop the application...")
    print("📊 Monitoring services for stability...\n")
    
    # Track request counts to detect stalled processes
    request_count = 0
    last_check_time = time.time()
    consecutive_failures = 0
    
    try:
        # Keep the script running until interrupted
        while True:
            # Check for crashed processes
            api_status = api_process.poll()
            streamlit_status = streamlit_process.poll()
            
            if api_status is not None:
                print(f"⚠️ API server stopped unexpectedly (exit code: {api_status})")
                # Restart the API server
                api_process = start_api_server()
                time.sleep(5)  # Give it time to start
                
            if streamlit_status is not None:
                print(f"⚠️ Streamlit app stopped unexpectedly (exit code: {streamlit_status})")
                # Restart Streamlit
                streamlit_process = start_streamlit_app()
                time.sleep(5)  # Give it time to start
            
            # More frequent health checks (every 10 seconds)
            if time.time() - last_check_time > 10:
                request_count += 1
                # Check API health
                if check_api_health():
                    print(f"✅ Health check #{request_count}: API is responsive")
                    consecutive_failures = 0  # Reset failure counter on success
                else:
                    consecutive_failures += 1
                    print(f"⚠️ Health check #{request_count}: API not responding (failure #{consecutive_failures})")
                    
                    # If we have 3 consecutive failures, restart the API
                    if consecutive_failures >= 3:
                        print("❌ API has been unresponsive for too long. Restarting...")
                        # Force terminate the API process
                        if api_process and api_process.poll() is None:
                            api_process.terminate()
                            time.sleep(1)
                            if api_process.poll() is None:
                                if sys.platform != 'win32':
                                    os.kill(api_process.pid, signal.SIGKILL)
                                else:
                                    subprocess.call(['taskkill', '/F', '/T', '/PID', str(api_process.pid)])
                        
                        # Start a new API process
                        api_process = start_api_server()
                        time.sleep(5)  # Give it time to start
                        consecutive_failures = 0  # Reset failure counter
                
                last_check_time = time.time()
            
            # Check memory usage of processes
            try:
                api_proc = psutil.Process(api_process.pid)
                api_memory = api_proc.memory_info().rss / (1024 * 1024)  # Convert to MB
                
                # If API is using too much memory (>500MB), restart it
                if api_memory > 500:
                    print(f"⚠️ API server using excessive memory: {api_memory:.2f} MB. Restarting...")
                    # Restart the API
                    if api_process and api_process.poll() is None:
                        api_process.terminate()
                        time.sleep(1)
                    api_process = start_api_server()
                    time.sleep(5)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
                
            # Sleep for a bit before checking again
            time.sleep(2)
            
    except KeyboardInterrupt:
        # Handle termination via Ctrl+C
        print("\n👋 Shutting down...")
    except Exception as e:
        # Handle any other exceptions
        print(f"\n❌ Error: {str(e)}")
    finally:
        shutdown_services(api_process, streamlit_process)

def shutdown_services(api_process, streamlit_process):
    """Safely shut down all running services"""
    # Clean up processes
    if api_process and api_process.poll() is None:
        print("Stopping API server...")
        try:
            api_process.terminate()
            # Give it a moment to terminate gracefully
            time.sleep(1)
            # Force kill if still running
            if api_process.poll() is None:
                if sys.platform != 'win32':
                    os.kill(api_process.pid, signal.SIGKILL)
                else:
                    subprocess.call(['taskkill', '/F', '/T', '/PID', str(api_process.pid)])
        except:
            print("Could not terminate API server gracefully")
    
    if streamlit_process and streamlit_process.poll() is None:
        print("Stopping Streamlit app...")
        try:
            streamlit_process.terminate()
            # Give it a moment to terminate gracefully
            time.sleep(1)
            # Force kill if still running
            if streamlit_process.poll() is None:
                if sys.platform != 'win32':
                    os.kill(streamlit_process.pid, signal.SIGKILL)
                else:
                    subprocess.call(['taskkill', '/F', '/T', '/PID', str(streamlit_process.pid)])
        except:
            print("Could not terminate Streamlit app gracefully")
    
    # Also kill any related processes that might still be hanging
    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
        try:
            cmdline = str(proc.info['cmdline']).lower()
            if ('uvicorn' in cmdline and '8001' in cmdline) or ('streamlit' in cmdline and 'run' in cmdline):
                print(f"Killing leftover process: {proc.info['pid']}")
                psutil.Process(proc.info['pid']).kill()
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    
    print("\n✅ Application stopped")

if __name__ == "__main__":
    main()