docker_entrypoint.py

#!/usr/bin/env python3
"""
Docker entrypoint script that decides which component to run based on environment variables.
"""
import os
import sys
import subprocess
import time
from pathlib import Path

def run_command(cmd):
    print(f"Running command: {' '.join(cmd)}")
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    
    # Stream the output
    for line in process.stdout:
        sys.stdout.write(line)
        sys.stdout.flush()
    
    process.wait()
    return process.returncode

def check_required_files():
    """Check if all required files exist"""
    required_files = [
        "pyscout_api.py",
        "deepinfra_client.py",
        "proxy_finder.py",
        "db_helper.py",
        "hf_utils.py",
    ]
    
    for file in required_files:
        if not Path(file).exists():
            print(f"ERROR: Required file '{file}' not found!")
            return False
    
    return True

def wait_for_mongodb():
    """Wait for MongoDB to be available"""
    import time
    import pymongo
    
    mongo_uri = os.environ.get("MONGODB_URI")
    if not mongo_uri:
        print("MongoDB URI not found in environment variables, skipping connection check")
        return True
    
    max_attempts = 30
    for attempt in range(max_attempts):
        try:
            client = pymongo.MongoClient(mongo_uri, serverSelectionTimeoutMS=5000)
            client.admin.command('ping')  # Simple command to check connection
            print(f"MongoDB connection successful after {attempt+1} attempts")
            return True
        except Exception as e:
            print(f"Attempt {attempt+1}/{max_attempts}: MongoDB not yet available. Waiting... ({str(e)})")
            time.sleep(2)
    
    print("ERROR: Failed to connect to MongoDB after multiple attempts")
    return False

def main():
    """Main entry point for the Docker container"""
    if not check_required_files():
        sys.exit(1)
    
    # Determine which component to run based on environment variable
    mode = os.environ.get("PYSCOUT_MODE", "api").lower()
    
    if mode == "api":
        print("Starting PyScoutAI API server")
        wait_for_mongodb()
        cmd = ["python", "pyscout_api.py"]
        return run_command(cmd)
        
    elif mode == "ui":
        print("Starting Gradio UI")
        cmd = ["python", "app.py"]
        return run_command(cmd)
        
    elif mode == "all":
        print("Starting both API server and UI")
        # Start API server in background
        api_process = subprocess.Popen(["python", "pyscout_api.py"])
        time.sleep(5)  # Wait for API to start
        
        # Start UI in foreground
        ui_cmd = ["python", "app.py"]
        ui_code = run_command(ui_cmd)
        
        # Kill API process when UI exits
        api_process.terminate()
        return ui_code
        
    else:
        print(f"ERROR: Unknown mode '{mode}'. Valid options: api, ui, all")
        return 1

if __name__ == "__main__":
    sys.exit(main())