simulation/simulation_worker.py

"""
Standalone Ray Simulation Worker

This module runs as a SEPARATE PROCESS from Celery, listening to Redis
for simulation requests and executing them using Ray.

Architecture:
  Celery (VLM) → publishes to Redis → Simulation Worker (Ray) → publishes results

Usage:
  python simulation/simulation_worker.py

This avoids the Ray/Celery conflict by keeping them in separate processes.
"""
import os
import sys
import json
import time
from redis.exceptions import ConnectionError as RedisConnectionError
import asyncio
import logging
import signal
from pathlib import Path
from datetime import datetime
from typing import Optional, List, Dict

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# Load environment variables
from dotenv import load_dotenv
load_dotenv(project_root / ".env")

import redis

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(name)s: %(message)s'
)
logger = logging.getLogger('simulation_worker')


class SimulationWorker:
    """
    Standalone worker that:
    1. Listens to Redis for simulation requests
    2. Runs simulations using Ray (via async orchestrator)
    3. Publishes results back to Redis
    """

    def __init__(self):
        """Initialize the worker"""
        self.redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
        self.redis_client = None
        self.running = False
        self.request_channel = "simulation_requests"
        self.reconnect_delay = 2
        self.max_reconnect_delay = 30

        # Database connection string
        self.db_url = os.getenv(
            "DATABASE_URL",
            "postgresql://agentsociety:dev_password@localhost:5432/agentsociety_db"
        )

        logger.info(f"SimulationWorker initialized with Redis: {self.redis_url}")

    def _create_redis_client(self):
        """Create a fresh Redis client, with TLS support for rediss:// URLs."""
        if self.redis_url.startswith("rediss://"):
            import ssl as ssl_module
            return redis.from_url(
                self.redis_url,
                ssl_cert_reqs=ssl_module.CERT_REQUIRED,
                socket_keepalive=True,
                health_check_interval=30,
            )
        return redis.from_url(
            self.redis_url,
            socket_keepalive=True,
            health_check_interval=30,
        )

    def run(self):
        """Main loop: listen for simulation requests with auto-reconnect."""
        self.running = True

        logger.info("=" * 60)
        logger.info("Ray Simulation Worker Started")
        logger.info(f"Listening for simulation requests on '{self.request_channel}'")
        logger.info("=" * 60)

        backoff = self.reconnect_delay

        while self.running:
            pubsub = None
            client = None

            try:
                client = self._create_redis_client()
                pubsub = client.pubsub(ignore_subscribe_messages=True)
                pubsub.subscribe(self.request_channel)

                self.redis_client = client
                backoff = self.reconnect_delay  # reset on successful connect

                logger.info("Connected to Redis and subscribed successfully")

                while self.running:
                    message = pubsub.get_message(timeout=1.0)
                    if message is None:
                        continue

                    if message["type"] != "message":
                        continue

                    try:
                        data = json.loads(message["data"])
                        self._handle_request(data)
                    except json.JSONDecodeError as e:
                        logger.error(f"Invalid JSON in message: {e}")
                    except Exception as e:
                        logger.exception(f"Error handling request: {e}")

            except (RedisConnectionError, OSError) as e:
                if not self.running:
                    break
                logger.warning(
                    f"Redis connection lost: {e}. Reconnecting in {backoff}s..."
                )
                time.sleep(backoff)
                backoff = min(backoff * 2, self.max_reconnect_delay)

            except Exception as e:
                if not self.running:
                    break
                logger.exception(
                    f"Unexpected worker loop error: {e}. Retrying in {backoff}s..."
                )
                time.sleep(backoff)
                backoff = min(backoff * 2, self.max_reconnect_delay)

            finally:
                try:
                    if pubsub is not None:
                        pubsub.close()
                except Exception:
                    pass

                try:
                    if client is not None:
                        client.close()
                except Exception:
                    pass

                if self.redis_client is client:
                    self.redis_client = None

    def _handle_request(self, request: dict):
        """Handle a simulation request"""
        simulation_id = request.get('simulation_id')

        if not simulation_id:
            logger.error("Request missing simulation_id")
            return

        logger.info(f"=== Starting Simulation: {simulation_id} ===")
        logger.info(f"  Agents: {request.get('num_agents', 10)}")
        logger.info(f"  Days: {request.get('simulation_days', 5)}")

        try:
            # Run the simulation using async orchestrator
            results = self._run_simulation(
                simulation_id=simulation_id,
                ad_content=request.get('ad_content', ''),
                demographic_filter=request.get('demographic_filter'),
                num_agents=request.get('num_agents', 10),
                simulation_days=request.get('simulation_days', 5),
                custom_agent_profiles=request.get('custom_agent_profiles'),
                use_custom_agents_only=request.get('use_custom_agents_only', False)
            )

            # Publish results to Redis
            result_message = {
                'simulation_id': simulation_id,
                'status': 'COMPLETED',
                'results': results,
                'completed_at': datetime.utcnow().isoformat()
            }

            if self.redis_client is not None:
                self.redis_client.publish(
                    "simulation_results",
                    json.dumps(result_message, default=str)
                )
            else:
                logger.warning(
                    f"Redis client unavailable; could not publish COMPLETED result "
                    f"for simulation {simulation_id}"
                )

            logger.info(f"=== Simulation Complete: {simulation_id} ===")
            logger.info(
                f"  Engagement Score: {results.get('engagement_score', 0):.2f}"
            )
            logger.info(
                f"  Risk Flags: {len(results.get('risk_flags', []))}"
            )

        except Exception as e:
            logger.error(f"Simulation {simulation_id} failed: {e}")

            # Publish error to Redis
            error_message = {
                'simulation_id': simulation_id,
                'status': 'FAILED',
                'error': str(e),
                'completed_at': datetime.utcnow().isoformat()
            }

            if self.redis_client is not None:
                self.redis_client.publish(
                    "simulation_results",
                    json.dumps(error_message)
                )
            else:
                logger.warning(
                    f"Redis client unavailable; could not publish FAILED result "
                    f"for simulation {simulation_id}"
                )

    def _run_simulation(
        self,
        simulation_id: str,
        ad_content: str,
        demographic_filter: Optional[dict],
        num_agents: int,
        simulation_days: int,
        custom_agent_profiles: Optional[List[Dict]] = None,
        use_custom_agents_only: bool = False
    ) -> dict:
        """Run the actual simulation using async orchestrator"""
        from simulation.run_simulation import run_simulation

        return run_simulation(
            experiment_id=simulation_id,
            ad_content=ad_content,
            demographic_filter=demographic_filter,
            num_agents=num_agents,
            simulation_days=simulation_days,
            redis_client=self.redis_client,
            custom_agent_profiles=custom_agent_profiles,
            use_custom_agents_only=use_custom_agents_only
        )

    def stop(self):
        """Stop the worker"""
        self.running = False
        logger.info("Stopping simulation worker...")


def main():
    """Entry point"""
    worker = SimulationWorker()

    # Handle graceful shutdown
    def signal_handler(sig, frame):
        print("\nShutting down gracefully...")
        worker.stop()
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # Run the worker
    worker.run()


if __name__ == "__main__":
    main()