local_quantum_simulator/Dockerfile_build/source/modules/rabbitmq.py

import asyncio
import json
from multiprocessing import Pipe, Process
from multiprocessing.connection import Connection
from time import sleep

import aio_pika
import pennylane
from aio_pika.abc import AbstractIncomingMessage
from connections.keycloak import (
    get_valid_access_token,
)
from connections.local_files import load_client_info
from connections.quantum_backend import (
    get_device_by_id,
)
from connections.rabbitmq import rabbitmq_manager
from fastcore.xtras import datetime
from modules.vqe import prepare_data, run_vqe
from pennylane.devices import Device
from uvicorn.main import logger

HEARTBEAT_EXCHANGE = "heartbeat"
HEARTBEAT_INTERVAL = 5
EXCHANGE_NAME = "progress_report"

# Global flag to track if VQE is busy
vqe_busy = False
vqe_busy_lock = asyncio.Lock()


# --- RabbitMQ Consumer Logic ---
async def process_message(message: AbstractIncomingMessage):
    """Your business logic for handling a message."""
    global vqe_busy

    async with message.process():
        body = message.body.decode()
        print(f"Received and processing: {body}")

        # Parse the JSON message
        data = json.loads(body)
        task_id = data.get("task_id")
        qubits_needed = data.get("qubits_needed")
        message_data = data.get("data", {})

        loop = asyncio.get_running_loop()

        # Extract molecular information from the text
        dev = pennylane.device("lightning.qubit", wires=qubits_needed)

        sleep(1)

        await publish_message(task_id, "", "IN SYSTEM")

        data = prepare_data(message_data)

        parent_conn, child_conn = Pipe()

        # Set busy flag before starting VQE
        async with vqe_busy_lock:
            vqe_busy = True

        try:
            # Run VQE in a separate process - will exit automatically when done
            vqe_process = Process(
                target=run_vqe_process,
                args=(child_conn, dev, data),
            )
            vqe_process.start()

            # Monitor pipe in a separate thread to avoid blocking the event loop
            def monitor_pipe_thread():
                while vqe_process.is_alive():
                    if parent_conn.poll():
                        logger.info("polling success")
                        result = parent_conn.recv()

                        if isinstance(result, dict) and result.get("status") == "ERROR":
                            asyncio.run_coroutine_threadsafe(
                                publish_message(task_id, json.dumps(result), "ERROR"),
                                loop,
                            )
                            break

                        # Check for sentinel
                        if result is None:
                            asyncio.run_coroutine_threadsafe(
                                publish_message(task_id, "", "COMPLETE"), loop
                            )
                            break

                        # Drain to get latest result
                        while parent_conn.poll():
                            result = parent_conn.recv()
                            if result is None:
                                asyncio.run_coroutine_threadsafe(
                                    publish_message(task_id, "", "COMPLETE"), loop
                                )
                                break

                        if result is not None:
                            logger.info(f"publishing {result}")
                            # Run async publish in the event loop
                            asyncio.run_coroutine_threadsafe(
                                publish_message(task_id, json.dumps(result)), loop
                            )
                    else:
                        import time

                        time.sleep(5)
                else:
                    if vqe_process.exitcode == 0:
                        asyncio.run_coroutine_threadsafe(
                            publish_message(task_id, "", "COMPLETE"), loop
                        )
                    else:
                        if parent_conn.poll():
                            logger.info("polling success")
                            result = parent_conn.recv()

                            if (
                                isinstance(result, dict)
                                and result.get("status") == "ERROR"
                            ):
                                print(f"ERROR ERROR ERROR ERROR ERROR ERROR")
                                asyncio.run_coroutine_threadsafe(
                                    publish_message(
                                        task_id, json.dumps(result), "ERROR"
                                    ),
                                    loop,
                                )

            # Run the monitor in a thread
            import threading

            monitor_thread = threading.Thread(target=monitor_pipe_thread, daemon=True)
            monitor_thread.start()

            # Wait for the process to finish
            await asyncio.to_thread(vqe_process.join)

            # Wait for monitor thread to finish
            monitor_thread.join(timeout=1)

            print(f"Completed task {task_id}")

        finally:
            # Clear busy flag after VQE completes (success or failure)
            async with vqe_busy_lock:
                vqe_busy = False


def run_vqe_process(conn: Connection, dev: Device, data: dict):
    """Run VQE in a separate process - process exits when this function returns."""
    try:
        # Run VQE
        run_vqe(conn=conn, dev1=dev, data=data)

        # Send sentinel to indicate completion
        conn.send(None)  # Signal completion
        conn.close()

    except Exception as e:
        logger.info(f"Error in VQE: {e}")
        import traceback

        # Send error information through the pipe BEFORE sending None
        error_msg = {
            "error": str(e),
            "traceback": traceback.format_exc(),
            "status": "ERROR",
        }
        conn.send(error_msg)  # Send error details
        conn.send(None)  # Signal completion
        import time

        time.sleep(5)
        conn.close()


async def consume_messages_topic():
    """
    Subscribe to team-specific qubit queues.
    Multiple systems subscribe to the SAME queue: team_{team_id}.qubits_{N}
    where N is the number of qubits this system can handle (1 to max_qubits).

    This enables round-robin task distribution among systems with sufficient qubits.
    """
    channel = await rabbitmq_manager.get_consumer_channel()
    if not channel:
        raise Exception("Failed to get consumer channel")

    client_info = load_client_info()
    if not client_info:
        raise Exception("Failed to load client info")

    system_id = client_info["system_id"]
    access_token = get_valid_access_token()
    device_data = get_device_by_id(system_id, access_token)
    teams = device_data["teams"]
    if not teams:
        print(f"System {system_id} is not part of any team. No queues to subscribe.")

    # Subscribe to qubit-specific queues for each team
    subscription_count = 0
    for team in teams:
        for qubits in range(1, team["num_qubits"] + 1):
            # Queue name format: team_{team_id}.qubits_{qubits}
            queue_name = f"team_{team['team']['team_id']}.qubits_{qubits}"

            # Declare the queue (durable, shared among multiple consumers)
            queue = await channel.declare_queue(
                queue_name,
                durable=True,
                arguments={
                    "x-max-priority": 100,  # Allow priorities 0-10
                },
            )

            # Start consuming from this queue
            await queue.consume(
                process_message,
                arguments={"x-priority": qubits, "x-priority-max": 10},
            )
            subscription_count += 1

            print(f"System {system_id} subscribed to queue: {queue_name}")

    print(f"System {system_id} subscribed to {subscription_count} queues")
    await asyncio.Future()  # Keep running


# --- Heartbeat Publisher Logic ---
async def publish_heartbeat():
    """Continuously publish 'alive' or 'busy' messages using the shared connection."""
    global vqe_busy

    channel = await rabbitmq_manager.get_heartbeat_channel()
    if not channel:
        raise Exception("Failed to get heartbeat channel")

    exchange = await channel.declare_exchange(
        HEARTBEAT_EXCHANGE, type=aio_pika.ExchangeType.FANOUT, durable=True
    )
    print(
        f"Heartbeat publisher started, sending status every {HEARTBEAT_INTERVAL} seconds..."
    )

    while True:
        try:
            client_info = load_client_info()
            if not client_info:
                raise Exception("Failed to load client info")

            system_id = client_info["system_id"]

            # Check current busy status
            async with vqe_busy_lock:
                current_status = "BUSY" if vqe_busy else "ONLINE"

            # Add timestamp for better monitoring
            heartbeat_message = {
                "device_id": system_id,
                "status": current_status,
                "timestamp": datetime.now().isoformat(),
            }

            message_body = json.dumps(heartbeat_message)
            await exchange.publish(
                aio_pika.Message(
                    body=message_body.encode(),
                    delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
                ),
                routing_key="",
            )
            print(
                f"Heartbeat sent: {current_status} at {asyncio.get_event_loop().time():.2f}"
            )
            await asyncio.sleep(HEARTBEAT_INTERVAL)

        except Exception as e:
            print(f"Error publishing heartbeat: {e}")
            await asyncio.sleep(1)


async def publish_message(task_id: int, message: str, status: str = "PROCESSING"):
    """Publish a message to the exchange using shared connection."""
    try:
        channel = await rabbitmq_manager.get_publisher_channel()
        if not channel:
            raise Exception("Failed to get publisher channel")

        # Declare exchange instead of queue
        exchange = await channel.declare_exchange(
            EXCHANGE_NAME,  # Using QUEUE_NAME as exchange name
            type=aio_pika.ExchangeType.DIRECT,  # or TOPIC/FANOUT based on your needs
            durable=True,
        )

        client_info = load_client_info()
        if not client_info:
            raise Exception("Failed to load client info")

        system_id = client_info["system_id"]

        # Publish directly to the exchange
        await exchange.publish(
            aio_pika.Message(
                body=message.encode(),
                headers={
                    "task_id": str(task_id),
                    "status": status,  # "PROCESSING" or "COMPLETE"
                    "system_id": str(system_id),
                },
                content_type="application/json",
                delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
            ),
            routing_key="",  # Empty routing key for direct exchange, or use task_id as routing key
        )
        print(f"Message published to exchange {EXCHANGE_NAME}: {message}")
    except Exception as e:
        print(f"Error publishing message: {e}")
        import traceback

        traceback.print_exc()