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_computation
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_computation(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()


# Global variable to track current subscriptions
_current_subscriptions = {}  # Format: {queue_name: {"team_id": team_id, "qubits": qubits, "consumer": consumer, "queue": queue}}


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.

    Dynamically manages subscriptions: adds new queues and removes obsolete ones.
    """
    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.")
        # Unsubscribe from all existing queues if no teams
        await _unsubscribe_all()
        return

    # Build expected queues from current teams
    expected_queues = set()
    for team in teams:
        for qubits in range(1, team["num_qubits"] + 1):
            queue_name = f"team_{team['team']['team_id']}.qubits_{qubits}"
            expected_queues.add(queue_name)

    # Get current subscriptions
    current_queues = set(_current_subscriptions.keys())

    # Queues to add (in expected but not current)
    queues_to_add = expected_queues - current_queues

    # Queues to remove (in current but not expected)
    queues_to_remove = current_queues - expected_queues

    # Unsubscribe from queues that no longer exist
    for queue_name in queues_to_remove:
        subscription_info = _current_subscriptions[queue_name]
        queue = subscription_info["queue"]
        cons = subscription_info["consumer"]
        try:
            # Cancel the consumer directly on the consumer object
            await queue.cancel(consumer_tag=cons)
            print(f"System {system_id} unsubscribed from queue: {queue_name}")
            del _current_subscriptions[queue_name]
        except Exception as e:
            print(f"Error unsubscribing from {queue_name}: {e}")

    # Subscribe to new queues
    for queue_name in queues_to_add:
        # Extract qubits from queue name
        qubits = int(queue_name.split(".qubits_")[1])
        team_id = queue_name.split(".qubits_")[0].replace("team_", "")

        # 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
        # In aio-pika, queue.consume returns a Consumer object
        consumer = await queue.consume(
            process_message,
            arguments={"x-priority": qubits, "x-priority-max": 10},
        )

        # Store subscription info with consumer object
        _current_subscriptions[queue_name] = {
            "team_id": team_id,
            "qubits": qubits,
            "consumer": consumer,
            "queue": queue,
        }

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

    # Print summary
    if queues_to_add or queues_to_remove:
        print(
            f"System {system_id} subscription update - Active: {len(_current_subscriptions)} queues, "
            f"Added: {len(queues_to_add)}, Removed: {len(queues_to_remove)}"
        )
    else:
        print(
            f"System {system_id} subscriptions unchanged - Active: {len(_current_subscriptions)} queues"
        )


async def _unsubscribe_all():
    """
    Helper function to unsubscribe from all queues.
    """
    for queue_name, subscription_info in list(_current_subscriptions.items()):
        consumer = subscription_info["consumer"]
        try:
            # Cancel the consumer directly on the consumer object
            await consumer.cancel()
            print(f"Unsubscribed from queue: {queue_name}")
            del _current_subscriptions[queue_name]
        except Exception as e:
            print(f"Error unsubscribing from {queue_name}: {e}")


def get_current_subscriptions():
    """
    Helper function to get current subscriptions for debugging.
    """
    # Return a serializable version without the consumer and queue objects
    return {
        queue_name: {
            "team_id": info["team_id"],
            "qubits": info["qubits"],
            "consumer_tag": info["consumer"].consumer_tag if info["consumer"] else None,
        }
        for queue_name, info in _current_subscriptions.items()
    }


# --- 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()