v0.1.0
Changes: - added rabbitmq hearbeat, consumer and publisher - fixed vqe to run with rabbitmq - added frontend via fasthtml - added proper .env configuration
This commit is contained in:
326
Dockerfile_build/source/modules/rabbitmq.py
Normal file
326
Dockerfile_build/source/modules/rabbitmq.py
Normal file
@@ -0,0 +1,326 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user