start_celery_worker.sh

#!/bin/bash

"""
Celery Worker Startup Script for Super-Agentic
==============================================

This script starts Celery workers with proper configuration for:
- Multiple priority queues (agents, messages, checkpoints, monitoring)
- Auto-scaling based on system load
- Health monitoring and graceful shutdown
- Environment validation
- Logging configuration

Usage:
    ./start_celery_worker.sh [environment] [worker_type]

Arguments:
    environment: development|staging|production (default: development)
    worker_type: all|agents|messages|monitoring (default: all)

Examples:
    ./start_celery_worker.sh                    # Start all workers in development
    ./start_celery_worker.sh production agents  # Start only agent workers in production
    ./start_celery_worker.sh staging all        # Start all workers in staging
"""

set -e  # Exit on any error

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Default configuration
ENVIRONMENT=${1:-development}
WORKER_TYPE=${2:-all}
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
API_DIR="$SCRIPT_DIR/apps/api"

# Worker configuration based on environment
case $ENVIRONMENT in
    "development")
        CONCURRENCY=2
        LOG_LEVEL="info"
        PREFETCH_MULTIPLIER=1
        MAX_TASKS_PER_CHILD=100
        ;;
    "staging")
        CONCURRENCY=4
        LOG_LEVEL="info"
        PREFETCH_MULTIPLIER=2
        MAX_TASKS_PER_CHILD=500
        ;;
    "production")
        CONCURRENCY=8
        LOG_LEVEL="warning"
        PREFETCH_MULTIPLIER=4
        MAX_TASKS_PER_CHILD=1000
        ;;
    *)
        echo -e "${RED}❌ Invalid environment: $ENVIRONMENT${NC}"
        echo "Valid environments: development, staging, production"
        exit 1
        ;;
esac

echo -e "${BLUE}🚀 Starting Celery Workers for Super-Agentic${NC}"
echo -e "${BLUE}Environment: $ENVIRONMENT${NC}"
echo -e "${BLUE}Worker Type: $WORKER_TYPE${NC}"
echo -e "${BLUE}Concurrency: $CONCURRENCY${NC}"
echo "================================================"

# Function to print status messages
print_status() {
    echo -e "${GREEN}✅ $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

print_error() {
    echo -e "${RED}❌ $1${NC}"
}

# Function to check if a command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to check environment variables
check_environment() {
    print_status "Checking environment configuration..."
    
    # Check if we're in the right directory
    if [ ! -f "$API_DIR/tasks/celery_app.py" ]; then
        print_error "Celery app not found. Please run from the super-agentic root directory."
        exit 1
    fi
    
    # Check if .env file exists
    if [ ! -f "$SCRIPT_DIR/.env" ]; then
        print_warning "No .env file found. Using environment variables or defaults."
    fi
    
    # Check required environment variables
    required_vars=("CELERY_BROKER_URL" "DATABASE_URL")
    missing_vars=()
    
    for var in "${required_vars[@]}"; do
        if [ -z "${!var}" ]; then
            missing_vars+=("$var")
        fi
    done
    
    if [ ${#missing_vars[@]} -ne 0 ]; then
        print_error "Missing required environment variables:"
        for var in "${missing_vars[@]}"; do
            echo "  - $var"
        done
        print_error "Please check your .env file or environment configuration."
        exit 1
    fi
    
    print_status "Environment configuration verified"
}

# Function to check system dependencies
check_dependencies() {
    print_status "Checking system dependencies..."
    
    # Check if Redis is running (for Celery broker)
    if command_exists redis-cli; then
        if ! redis-cli ping >/dev/null 2>&1; then
            print_warning "Redis server not responding. Please ensure Redis is running."
            print_warning "To start Redis: brew services start redis (macOS) or sudo systemctl start redis (Linux)"
        else
            print_status "Redis connection verified"
        fi
    else
        print_warning "redis-cli not found. Cannot verify Redis connection."
    fi
    
    # Check if PostgreSQL is accessible
    if command_exists psql; then
        # Extract database info from DATABASE_URL if available
        if [ -n "$DATABASE_URL" ]; then
            if psql "$DATABASE_URL" -c "SELECT 1;" >/dev/null 2>&1; then
                print_status "Database connection verified"
            else
                print_warning "Database connection failed. Please check DATABASE_URL."
            fi
        fi
    else
        print_warning "psql not found. Cannot verify database connection."
    fi
}

# Function to activate virtual environment
activate_virtualenv() {
    print_status "Activating virtual environment..."
    
    # Look for virtual environment in common locations
    venv_paths=(
        "$SCRIPT_DIR/venv"
        "$SCRIPT_DIR/.venv"
        "$API_DIR/venv"
        "$API_DIR/.venv"
    )
    
    for venv_path in "${venv_paths[@]}"; do
        if [ -f "$venv_path/bin/activate" ]; then
            source "$venv_path/bin/activate"
            print_status "Virtual environment activated: $venv_path"
            return 0
        fi
    done
    
    print_warning "No virtual environment found. Using system Python."
    print_warning "Consider creating a virtual environment: python -m venv venv"
}

# Function to install dependencies
check_python_dependencies() {
    print_status "Checking Python dependencies..."
    
    cd "$API_DIR"
    
    # Check if required packages are installed
    required_packages=("celery" "redis" "psutil" "langchain" "sqlalchemy")
    missing_packages=()
    
    for package in "${required_packages[@]}"; do
        if ! python -c "import $package" >/dev/null 2>&1; then
            missing_packages+=("$package")
        fi
    done
    
    if [ ${#missing_packages[@]} -ne 0 ]; then
        print_warning "Missing Python packages: ${missing_packages[*]}"
        print_warning "Installing missing packages..."
        pip install "${missing_packages[@]}"
    fi
    
    print_status "Python dependencies verified"
}

# Function to create log directory
setup_logging() {
    print_status "Setting up logging..."
    
    LOG_DIR="$SCRIPT_DIR/logs"
    mkdir -p "$LOG_DIR"
    
    # Create log files if they don't exist
    touch "$LOG_DIR/celery_worker.log"
    touch "$LOG_DIR/celery_beat.log"
    
    print_status "Log directory created: $LOG_DIR"
}

# Function to start worker based on type
start_worker() {
    local worker_type=$1
    local worker_name="super-agentic-$worker_type-worker"
    local log_file="$SCRIPT_DIR/logs/celery_${worker_type}_worker.log"
    
    print_status "Starting $worker_type worker..."
    
    cd "$API_DIR"
    
    case $worker_type in
        "all")
            # Start worker for all queues
            celery -A tasks.celery_app worker \
                --hostname="$worker_name@%h" \
                --queues=agents,messages,checkpoints,monitoring \
                --concurrency=$CONCURRENCY \
                --loglevel=$LOG_LEVEL \
                --logfile="$log_file" \
                --pidfile="$SCRIPT_DIR/logs/celery_worker.pid" \
                --prefetch-multiplier=$PREFETCH_MULTIPLIER \
                --max-tasks-per-child=$MAX_TASKS_PER_CHILD \
                --detach
            ;;
        "agents")
            # Start worker for agent tasks only (high priority)
            celery -A tasks.celery_app worker \
                --hostname="$worker_name@%h" \
                --queues=agents \
                --concurrency=$((CONCURRENCY * 2)) \
                --loglevel=$LOG_LEVEL \
                --logfile="$log_file" \
                --pidfile="$SCRIPT_DIR/logs/celery_agents_worker.pid" \
                --prefetch-multiplier=$PREFETCH_MULTIPLIER \
                --max-tasks-per-child=$MAX_TASKS_PER_CHILD \
                --detach
            ;;
        "messages")
            # Start worker for messaging tasks
            celery -A tasks.celery_app worker \
                --hostname="$worker_name@%h" \
                --queues=messages \
                --concurrency=$CONCURRENCY \
                --loglevel=$LOG_LEVEL \
                --logfile="$log_file" \
                --pidfile="$SCRIPT_DIR/logs/celery_messages_worker.pid" \
                --prefetch-multiplier=$PREFETCH_MULTIPLIER \
                --max-tasks-per-child=$MAX_TASKS_PER_CHILD \
                --detach
            ;;
        "monitoring")
            # Start worker for monitoring tasks (low priority)
            celery -A tasks.celery_app worker \
                --hostname="$worker_name@%h" \
                --queues=monitoring,checkpoints \
                --concurrency=$((CONCURRENCY / 2)) \
                --loglevel=$LOG_LEVEL \
                --logfile="$log_file" \
                --pidfile="$SCRIPT_DIR/logs/celery_monitoring_worker.pid" \
                --prefetch-multiplier=1 \
                --max-tasks-per-child=$MAX_TASKS_PER_CHILD \
                --detach
            ;;
        *)
            print_error "Invalid worker type: $worker_type"
            echo "Valid worker types: all, agents, messages, monitoring"
            exit 1
            ;;
    esac
    
    print_status "$worker_type worker started successfully"
}

# Function to start Celery Beat (scheduler)
start_beat() {
    print_status "Starting Celery Beat scheduler..."
    
    cd "$API_DIR"
    
    celery -A tasks.celery_app beat \
        --loglevel=$LOG_LEVEL \
        --logfile="$SCRIPT_DIR/logs/celery_beat.log" \
        --pidfile="$SCRIPT_DIR/logs/celery_beat.pid" \
        --detach
    
    print_status "Celery Beat scheduler started successfully"
}

# Function to show worker status
show_status() {
    print_status "Checking worker status..."
    
    cd "$API_DIR"
    
    echo ""
    echo "Active Celery Workers:"
    celery -A tasks.celery_app inspect active 2>/dev/null || echo "No active workers found"
    
    echo ""
    echo "Celery Worker Stats:"
    celery -A tasks.celery_app inspect stats 2>/dev/null || echo "No worker stats available"
    
    echo ""
    echo "Running Processes:"
    ps aux | grep -E "(celery|super-agentic)" | grep -v grep || echo "No Celery processes found"
}

# Function to cleanup on exit
cleanup() {
    print_status "Cleaning up..."
    # Any cleanup tasks can be added here
}

# Set up signal handlers for graceful shutdown
trap cleanup EXIT INT TERM

# Main execution
main() {
    # Load environment variables if .env exists
    if [ -f "$SCRIPT_DIR/.env" ]; then
        set -a  # Automatically export all variables
        source "$SCRIPT_DIR/.env"
        set +a  # Turn off automatic export
    fi
    
    # Run all checks
    check_environment
    check_dependencies
    activate_virtualenv
    check_python_dependencies
    setup_logging
    
    # Start the requested worker type
    start_worker "$WORKER_TYPE"
    
    # Start Celery Beat scheduler (only for 'all' worker type in production)
    if [ "$WORKER_TYPE" = "all" ] && [ "$ENVIRONMENT" != "development" ]; then
        start_beat
    fi
    
    # Show status
    sleep 3  # Give workers time to start
    show_status
    
    echo ""
    print_status "Celery workers started successfully!"
    print_status "Logs are available in: $SCRIPT_DIR/logs/"
    print_status "To stop workers: make celery-stop"
    print_status "To monitor workers: make celery-monitor"
    
    echo ""
    echo -e "${BLUE}📊 Useful Commands:${NC}"
    echo "  View logs: tail -f $SCRIPT_DIR/logs/celery_worker.log"
    echo "  Monitor workers: celery -A tasks.celery_app inspect active"
    echo "  Stop workers: pkill -f 'celery worker'"
    echo "  Worker stats: celery -A tasks.celery_app inspect stats"
}

# Run main function
main "$@"