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