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:
86
Dockerfile_build/source/connections/keycloak.py
Normal file
86
Dockerfile_build/source/connections/keycloak.py
Normal file
@@ -0,0 +1,86 @@
|
||||
import os
|
||||
from datetime import timedelta
|
||||
|
||||
from connections.local_files import (
|
||||
delete_client_info,
|
||||
delete_token,
|
||||
delete_verifier,
|
||||
load_token,
|
||||
save_token,
|
||||
)
|
||||
from fastcore.xtras import datetime
|
||||
|
||||
from keycloak import KeycloakOpenID
|
||||
|
||||
KEYCLOAK_CONFIG = {
|
||||
"server_url": os.environ["KEYCLOAK_URL"],
|
||||
"realm_name": os.environ["KEYCLOAK_REALM_NAME"],
|
||||
"client_id": os.environ["KEACLOAK_CLIENT_ID"],
|
||||
}
|
||||
|
||||
# Инициализация клиента Keycloak
|
||||
keycloak_openid = KeycloakOpenID(
|
||||
server_url=KEYCLOAK_CONFIG["server_url"],
|
||||
realm_name=KEYCLOAK_CONFIG["realm_name"],
|
||||
client_id=KEYCLOAK_CONFIG["client_id"],
|
||||
)
|
||||
|
||||
|
||||
def refresh_access_token() -> dict | None:
|
||||
"""Refresh the access token using refresh token"""
|
||||
token_data = load_token()
|
||||
if not token_data or not token_data.get("refresh_token"):
|
||||
return None
|
||||
|
||||
try:
|
||||
# Use Keycloak's refresh token endpoint
|
||||
refresh_token = token_data["refresh_token"]
|
||||
|
||||
# Get new token pair using refresh token
|
||||
new_tokens = keycloak_openid.refresh_token(refresh_token)
|
||||
|
||||
# Calculate new expiration time
|
||||
expires_in = new_tokens.get("expires_in", 3600)
|
||||
expires_at = datetime.now() + timedelta(seconds=expires_in)
|
||||
|
||||
# Update token data
|
||||
updated_token_data = {
|
||||
"access_token": new_tokens["access_token"],
|
||||
"refresh_token": new_tokens.get(
|
||||
"refresh_token", refresh_token
|
||||
), # Some providers return new refresh token
|
||||
"expires_at": expires_at.isoformat(),
|
||||
"id_token": new_tokens.get("id_token", token_data.get("id_token")),
|
||||
"user_info": token_data.get("user_info"), # Preserve user info
|
||||
}
|
||||
|
||||
# Save updated tokens
|
||||
save_token(updated_token_data)
|
||||
|
||||
return updated_token_data
|
||||
|
||||
except Exception as e:
|
||||
# If refresh fails, clear all tokens and force re-login
|
||||
delete_token()
|
||||
delete_client_info()
|
||||
delete_verifier()
|
||||
return None
|
||||
|
||||
|
||||
def get_valid_access_token() -> str | None:
|
||||
"""Get a valid access token, refreshing if necessary"""
|
||||
token_data = load_token()
|
||||
|
||||
if not token_data or not token_data.get("expires_at"):
|
||||
return None
|
||||
|
||||
# Check if token is expired or about to expire (within 60 seconds)
|
||||
expires_at = datetime.fromisoformat(str(token_data.get("expires_at")))
|
||||
if expires_at <= datetime.now() + timedelta(seconds=60):
|
||||
# Token is expired or about to expire, try to refresh
|
||||
new_token_data = refresh_access_token()
|
||||
if new_token_data:
|
||||
return new_token_data["access_token"]
|
||||
return None
|
||||
|
||||
return token_data["access_token"]
|
||||
69
Dockerfile_build/source/connections/local_files.py
Normal file
69
Dockerfile_build/source/connections/local_files.py
Normal file
@@ -0,0 +1,69 @@
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Локальные файлы для хранения данных
|
||||
TOKEN_FILE = Path(os.environ["STORAGE_PATH"] + "/token.json")
|
||||
VERIFIER_FILE = Path(os.environ["STORAGE_PATH"] + "/pkce_verifier.json")
|
||||
CLIENT_INFO_FILE = Path(os.environ["STORAGE_PATH"] + "/client_info.json")
|
||||
|
||||
|
||||
def save_token(token_data: dict):
|
||||
"""Сохранить токен в файл"""
|
||||
with open(TOKEN_FILE, "w") as f:
|
||||
json.dump(token_data, f, indent=2)
|
||||
|
||||
|
||||
def load_token() -> dict | None:
|
||||
"""Загрузить токен из файла"""
|
||||
if not TOKEN_FILE.exists():
|
||||
return None
|
||||
with open(TOKEN_FILE, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def delete_token():
|
||||
"""Удалить файл с токеном"""
|
||||
if TOKEN_FILE.exists():
|
||||
TOKEN_FILE.unlink()
|
||||
|
||||
|
||||
def save_verifier(verifier: str):
|
||||
"""Сохранить PKCE верификатор в файл"""
|
||||
with open(VERIFIER_FILE, "w") as f:
|
||||
json.dump({"verifier": verifier}, f)
|
||||
|
||||
|
||||
def load_verifier() -> str | None:
|
||||
"""Загрузить PKCE верификатор из файла"""
|
||||
if not VERIFIER_FILE.exists():
|
||||
return None
|
||||
with open(VERIFIER_FILE, "r") as f:
|
||||
data = json.load(f)
|
||||
return data.get("verifier")
|
||||
|
||||
|
||||
def delete_verifier():
|
||||
"""Удалить файл с PKCE верификатором"""
|
||||
if VERIFIER_FILE.exists():
|
||||
VERIFIER_FILE.unlink()
|
||||
|
||||
|
||||
def save_client_info(client_info: dict):
|
||||
"""Сохранить информацию о клиенте локально"""
|
||||
with open(CLIENT_INFO_FILE, "w") as f:
|
||||
json.dump(client_info, f, indent=2)
|
||||
|
||||
|
||||
def load_client_info() -> dict | None:
|
||||
"""Загрузить информацию о клиенте"""
|
||||
if not CLIENT_INFO_FILE.exists():
|
||||
return None
|
||||
with open(CLIENT_INFO_FILE, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def delete_client_info():
|
||||
"""Удалить файл с информацией о клиенте"""
|
||||
if CLIENT_INFO_FILE.exists():
|
||||
CLIENT_INFO_FILE.unlink()
|
||||
139
Dockerfile_build/source/connections/quantum_backend.py
Normal file
139
Dockerfile_build/source/connections/quantum_backend.py
Normal file
@@ -0,0 +1,139 @@
|
||||
import os
|
||||
from typing import Any, Dict
|
||||
|
||||
import requests
|
||||
|
||||
# Global var from env
|
||||
QUANTUM_BACKEND_URL = os.getenv(
|
||||
"QUANTUM_BACKEND_URL", os.environ["QUNATUM_BACKEND_URL"]
|
||||
)
|
||||
|
||||
|
||||
def get_or_create_device_by_name(
|
||||
system_name: str, max_qubits: int, access_token
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get or create a device by name.
|
||||
|
||||
Args:
|
||||
system_name: Name of the system
|
||||
max_qubits: Maximum number of qubits
|
||||
|
||||
Returns:
|
||||
Device data as dictionary
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If the request fails
|
||||
"""
|
||||
url = f"{QUANTUM_BACKEND_URL}/machine/"
|
||||
payload = {"system_name": system_name, "max_qubits": max_qubits}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.RequestException as e:
|
||||
raise Exception(f"Failed to get or create device: {e}") from e
|
||||
|
||||
|
||||
def get_device_by_id(system_id: int, access_token) -> Dict[str, Any]:
|
||||
"""
|
||||
Get device by ID.
|
||||
|
||||
Args:
|
||||
system_id: ID of the system
|
||||
|
||||
Returns:
|
||||
Device data as dictionary
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If the request fails
|
||||
"""
|
||||
url = f"{QUANTUM_BACKEND_URL}/machine/system"
|
||||
params = {"system_id": system_id}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.get(url, params=params, headers=headers)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.RequestException as e:
|
||||
raise Exception(f"Failed to get device by ID {system_id}: {e}") from e
|
||||
|
||||
|
||||
def update_device_data(
|
||||
system_id: int,
|
||||
system_name: str,
|
||||
system_description: str,
|
||||
max_qubits: int,
|
||||
access_token,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Update device data.
|
||||
|
||||
Args:
|
||||
system_id: ID of the system
|
||||
system_name: New system name
|
||||
system_description: New system description
|
||||
max_qubits: New maximum qubits
|
||||
|
||||
Returns:
|
||||
Updated device data as dictionary
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If the request fails
|
||||
"""
|
||||
url = f"{QUANTUM_BACKEND_URL}/machine"
|
||||
payload = {
|
||||
"system_id": system_id,
|
||||
"system_name": system_name,
|
||||
"system_description": system_description,
|
||||
"max_qubits": max_qubits,
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.put(url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.RequestException as e:
|
||||
raise Exception(f"Failed to update device data for ID {system_id}: {e}") from e
|
||||
|
||||
|
||||
def get_comp_system_file(experiment_type_id: int) -> bytes:
|
||||
"""
|
||||
Get comp_system file as a string/blob.
|
||||
|
||||
Args:
|
||||
experiment_type_id: ID of the experiment type
|
||||
|
||||
Returns:
|
||||
Python module content as string
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If the request fails
|
||||
"""
|
||||
url = f"{QUANTUM_BACKEND_URL}/experiment/types/comp-system"
|
||||
params = {"experiment_type_id": experiment_type_id}
|
||||
|
||||
try:
|
||||
response = requests.get(url, params=params)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
except requests.RequestException as e:
|
||||
raise Exception(
|
||||
f"Failed to get comp_system file for experiment_type_id {experiment_type_id}: {e}"
|
||||
) from e
|
||||
97
Dockerfile_build/source/connections/rabbitmq.py
Normal file
97
Dockerfile_build/source/connections/rabbitmq.py
Normal file
@@ -0,0 +1,97 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import aio_pika
|
||||
from aio_pika.abc import (
|
||||
AbstractChannel,
|
||||
AbstractRobustConnection,
|
||||
)
|
||||
from connections.keycloak import (
|
||||
get_valid_access_token,
|
||||
)
|
||||
from connections.local_files import (
|
||||
load_client_info,
|
||||
)
|
||||
|
||||
|
||||
class RabbitMQManager:
|
||||
"""Singleton manager for RabbitMQ connection and channels."""
|
||||
|
||||
_instance: Optional["RabbitMQManager"] = None
|
||||
_connection: Optional[AbstractRobustConnection] = None
|
||||
_consumer_channel: Optional[AbstractChannel] = None
|
||||
_publisher_channel: Optional[AbstractChannel] = None
|
||||
_heartbeat_channel: Optional[AbstractChannel] = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
async def connect(self):
|
||||
"""Establish the main connection if not already connected."""
|
||||
client_info = load_client_info()
|
||||
if not client_info:
|
||||
return None
|
||||
access_token = get_valid_access_token()
|
||||
if not access_token:
|
||||
return None
|
||||
if self._connection is None or self._connection.is_closed:
|
||||
print("Creating new RabbitMQ connection...")
|
||||
self._connection = await aio_pika.connect_robust(
|
||||
host=os.environ["RABBITMQ_HOST"],
|
||||
port=int(os.environ["RABBITMQ_PORT"]),
|
||||
login=str(client_info["system_id"]),
|
||||
password=access_token,
|
||||
virtualhost="/",
|
||||
)
|
||||
print("RabbitMQ connection established")
|
||||
return self._connection
|
||||
|
||||
async def get_consumer_channel(self):
|
||||
"""Get channel for consuming messages."""
|
||||
if self._connection:
|
||||
if self._consumer_channel is None or self._consumer_channel.is_closed:
|
||||
self._consumer_channel = await self._connection.channel()
|
||||
await self._consumer_channel.set_qos(prefetch_count=1, global_=True)
|
||||
print("Consumer channel created")
|
||||
return self._consumer_channel
|
||||
return None
|
||||
|
||||
async def get_publisher_channel(self):
|
||||
"""Get channel for publishing regular messages."""
|
||||
if self._connection:
|
||||
if self._publisher_channel is None or self._publisher_channel.is_closed:
|
||||
self._publisher_channel = await self._connection.channel()
|
||||
print("Publisher channel created")
|
||||
return self._publisher_channel
|
||||
return None
|
||||
|
||||
async def get_heartbeat_channel(self):
|
||||
"""Get channel for heartbeat publishing."""
|
||||
if self._connection:
|
||||
if self._heartbeat_channel is None or self._heartbeat_channel.is_closed:
|
||||
self._heartbeat_channel = await self._connection.channel()
|
||||
print("Heartbeat channel created")
|
||||
return self._heartbeat_channel
|
||||
return None
|
||||
|
||||
async def close(self):
|
||||
"""Close all channels and the main connection gracefully."""
|
||||
print("Closing RabbitMQ channels and connection...")
|
||||
|
||||
for channel in [
|
||||
self._consumer_channel,
|
||||
self._publisher_channel,
|
||||
self._heartbeat_channel,
|
||||
]:
|
||||
if channel and not channel.is_closed:
|
||||
await channel.close()
|
||||
|
||||
if self._connection and not self._connection.is_closed:
|
||||
await self._connection.close()
|
||||
print("RabbitMQ connection closed")
|
||||
|
||||
|
||||
# Create global singleton instance
|
||||
rabbitmq_manager = RabbitMQManager()
|
||||
72
Dockerfile_build/source/main.py
Normal file
72
Dockerfile_build/source/main.py
Normal file
@@ -0,0 +1,72 @@
|
||||
import asyncio
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import uvicorn
|
||||
from connections.rabbitmq import rabbitmq_manager
|
||||
from fasthtml.core import FastHTML, Mount
|
||||
from modules.fasthtml import app as fasthtml_app
|
||||
from modules.rabbitmq import consume_messages_topic, publish_heartbeat
|
||||
|
||||
|
||||
async def start_heartbeat():
|
||||
"""Start heartbeat with auto-reconnect"""
|
||||
while True:
|
||||
try:
|
||||
await publish_heartbeat()
|
||||
except Exception as e:
|
||||
print(f"Heartbeat failed: {e}")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
|
||||
async def start_consumer():
|
||||
"""Start consumer with auto-reconnect"""
|
||||
while True:
|
||||
try:
|
||||
await consume_messages_topic()
|
||||
except Exception as e:
|
||||
print(f"Consumer failed: {e}")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
|
||||
async def connect_with_retry():
|
||||
"""Retry RabbitMQ connection until successful"""
|
||||
while True:
|
||||
try:
|
||||
await rabbitmq_manager.connect()
|
||||
if (
|
||||
rabbitmq_manager._connection
|
||||
and not rabbitmq_manager._connection.is_closed
|
||||
):
|
||||
print("RabbitMQ connected")
|
||||
asyncio.create_task(start_heartbeat())
|
||||
asyncio.create_task(start_consumer())
|
||||
return
|
||||
else:
|
||||
await asyncio.sleep(5)
|
||||
except Exception as e:
|
||||
print(f"RabbitMQ connection failed: {e}, retrying in 5 seconds...")
|
||||
await asyncio.sleep(5)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app):
|
||||
print("Starting up...")
|
||||
|
||||
asyncio.create_task(connect_with_retry())
|
||||
|
||||
yield
|
||||
|
||||
print("Shutting down...")
|
||||
for task in asyncio.all_tasks():
|
||||
if task is not asyncio.current_task():
|
||||
task.cancel()
|
||||
|
||||
await rabbitmq_manager.close()
|
||||
|
||||
|
||||
app = FastHTML(routes=[Mount("", fasthtml_app, name="FastHTML")])
|
||||
app.set_lifespan(lifespan)
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run("main:app", host="0.0.0.0", port=int(os.environ["PORT"]), reload=True)
|
||||
505
Dockerfile_build/source/modules/fasthtml.py
Normal file
505
Dockerfile_build/source/modules/fasthtml.py
Normal file
@@ -0,0 +1,505 @@
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import psutil
|
||||
from connections.keycloak import get_valid_access_token, keycloak_openid
|
||||
from connections.local_files import (
|
||||
delete_client_info,
|
||||
delete_token,
|
||||
delete_verifier,
|
||||
load_client_info,
|
||||
load_token,
|
||||
load_verifier,
|
||||
save_client_info,
|
||||
save_token,
|
||||
save_verifier,
|
||||
)
|
||||
from connections.quantum_backend import get_device_by_id, get_or_create_device_by_name
|
||||
from connections.rabbitmq import (
|
||||
rabbitmq_manager,
|
||||
)
|
||||
from fasthtml.common import (
|
||||
H1,
|
||||
H2,
|
||||
H3,
|
||||
A,
|
||||
Button,
|
||||
Div,
|
||||
FastHTML,
|
||||
Form,
|
||||
Input,
|
||||
Meta,
|
||||
P,
|
||||
RedirectResponse,
|
||||
Script,
|
||||
Span,
|
||||
Strong,
|
||||
Titled,
|
||||
)
|
||||
from fasthtml.pico import Card, Container, picolink
|
||||
from fasthtml.xtend import Style
|
||||
from keycloak.pkce_utils import generate_code_challenge, generate_code_verifier
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
REDIRECT_URI = f"http://localhost:{os.environ['PORT']}/callback"
|
||||
|
||||
css = Style(":root {--pico-font-size:90%,--pico-font-family: Pacifico, cursive;}")
|
||||
app = FastHTML(hdrs=(picolink, css))
|
||||
|
||||
|
||||
@app.route("/", methods="get")
|
||||
async def main_page():
|
||||
# Проверяем наличие валидного токена
|
||||
token_data = load_token()
|
||||
client_info = load_client_info()
|
||||
|
||||
if token_data:
|
||||
expires_at = datetime.fromisoformat(token_data.get("expires_at", "2000-01-01"))
|
||||
if expires_at > datetime.now():
|
||||
# Если есть токен, но нет информации о клиенте, перенаправляем на регистрацию
|
||||
if not client_info:
|
||||
return RedirectResponse("/register-client", status_code=303)
|
||||
return RedirectResponse("/dashboard", status_code=303)
|
||||
else:
|
||||
delete_token()
|
||||
delete_client_info()
|
||||
|
||||
return Titled(
|
||||
"Главная",
|
||||
H1("Добро пожаловать в Квантовый Симулятор"),
|
||||
P("Пожалуйста, войдите для продолжения"),
|
||||
A("Войти через Keycloak", href="/login", cls="button"),
|
||||
)
|
||||
|
||||
|
||||
@app.route("/login", methods="get")
|
||||
async def login():
|
||||
# Генерация PKCE кода верификатора и запроса
|
||||
code_verifier = generate_code_verifier()
|
||||
code_challenge, code_challenge_method = generate_code_challenge(
|
||||
code_verifier, method="S256"
|
||||
)
|
||||
|
||||
# Сохраняем верификатор локально
|
||||
save_verifier(code_verifier)
|
||||
|
||||
# Формирование URL авторизации
|
||||
auth_url = keycloak_openid.auth_url(
|
||||
redirect_uri=REDIRECT_URI,
|
||||
scope="openid profile email",
|
||||
code_challenge=code_challenge,
|
||||
code_challenge_method=code_challenge_method,
|
||||
)
|
||||
|
||||
return RedirectResponse(auth_url, status_code=303)
|
||||
|
||||
|
||||
@app.route("/callback", methods="get")
|
||||
async def callback(request):
|
||||
code = request.query_params.get("code")
|
||||
|
||||
if not code:
|
||||
return Titled(
|
||||
"Ошибка", P("Код не найден"), A("Попробовать снова", href="/login")
|
||||
)
|
||||
|
||||
# Загружаем верификатор
|
||||
code_verifier = load_verifier()
|
||||
if not code_verifier:
|
||||
return Titled(
|
||||
"Ошибка",
|
||||
P("PKCE верификатор не найден"),
|
||||
A("Попробовать снова", href="/login"),
|
||||
)
|
||||
|
||||
# Очищаем верификатор сразу после использования
|
||||
delete_verifier()
|
||||
|
||||
try:
|
||||
# Обмен кода на токен
|
||||
token = keycloak_openid.token(
|
||||
grant_type="authorization_code",
|
||||
code=code,
|
||||
redirect_uri=REDIRECT_URI,
|
||||
code_verifier=code_verifier,
|
||||
)
|
||||
|
||||
# Расчет времени истечения
|
||||
expires_in = token.get("expires_in", 3600)
|
||||
expires_at = datetime.now() + timedelta(seconds=expires_in)
|
||||
|
||||
# Получение информации о пользователе
|
||||
access_token = token.get("access_token")
|
||||
if not access_token:
|
||||
raise Exception
|
||||
user_info = keycloak_openid.userinfo(access_token)
|
||||
|
||||
# Сохраняем токен
|
||||
token_data_to_save = {
|
||||
"access_token": access_token,
|
||||
"refresh_token": token.get("refresh_token"),
|
||||
"expires_at": expires_at.isoformat(),
|
||||
"id_token": token.get("id_token"),
|
||||
"user_info": user_info,
|
||||
}
|
||||
|
||||
save_token(token_data_to_save)
|
||||
|
||||
# Перенаправление на регистрацию клиента (проверит, зарегистрирован ли уже)
|
||||
return RedirectResponse("/register-client", status_code=303)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Ошибка при обмене токена: {e}")
|
||||
return Titled("Ошибка", P(str(e)), A("Попробовать снова", href="/login"))
|
||||
|
||||
|
||||
@app.route("/register-client", methods="get")
|
||||
async def register_client(request):
|
||||
"""Страница регистрации клиента в центральном микросервисе"""
|
||||
token_data = load_token()
|
||||
|
||||
if not token_data:
|
||||
return RedirectResponse("/login", status_code=303)
|
||||
|
||||
# Проверяем, не зарегистрирован ли уже
|
||||
client_info = load_client_info()
|
||||
if client_info:
|
||||
return RedirectResponse("/dashboard", status_code=303)
|
||||
|
||||
# Проверяем истечение токена
|
||||
expires_at = datetime.fromisoformat(str(token_data.get("expires_at")))
|
||||
if expires_at <= datetime.now():
|
||||
delete_token()
|
||||
return RedirectResponse("/login", status_code=303)
|
||||
|
||||
user_info = token_data.get("user_info", {})
|
||||
username = user_info.get(
|
||||
"preferred_username", user_info.get("name", "Пользователь")
|
||||
)
|
||||
|
||||
return Titled(
|
||||
"Регистрация клиента",
|
||||
Container(
|
||||
Card(
|
||||
H2("Добро пожаловать в Квантовый Симулятор!"),
|
||||
P(f"Здравствуйте, {username}!"),
|
||||
P("Это устройство должно быть зарегистрировано в центральном сервисе."),
|
||||
P(
|
||||
"Пожалуйста, укажите уникальное имя и количество кубит для этого клиента:"
|
||||
),
|
||||
Form(
|
||||
Input(
|
||||
type="text",
|
||||
name="client_name",
|
||||
id="client_name",
|
||||
placeholder="например: рабочий-ноутбук, домашний-пк, raspberry-pi-01",
|
||||
required=True,
|
||||
style="width: 100%; padding: 10px; margin: 10px 0;",
|
||||
),
|
||||
Input(
|
||||
type="number",
|
||||
name="number_qubits",
|
||||
id="number_qubits",
|
||||
placeholder=0,
|
||||
required=True,
|
||||
style="width: 100%; padding: 10px; margin: 10px 0;",
|
||||
),
|
||||
Div(
|
||||
Button(
|
||||
"Зарегистрировать клиент",
|
||||
type="submit",
|
||||
cls="button",
|
||||
style="background-color: #4CAF50;",
|
||||
),
|
||||
style="margin-top: 10px;",
|
||||
),
|
||||
hx_post="/register-client",
|
||||
hx_target="#registration-result",
|
||||
hx_swap="innerHTML",
|
||||
style="margin-top: 20px;",
|
||||
),
|
||||
Div(id="registration-result"),
|
||||
),
|
||||
Div(
|
||||
Button(
|
||||
"Выйти",
|
||||
onclick="window.location.href='/logout'",
|
||||
cls="button",
|
||||
style="background-color: #f44336; margin-top: 20px;",
|
||||
),
|
||||
style="text-align: center;",
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@app.route("/register-client", methods="post")
|
||||
async def register_client_post(request):
|
||||
"""Обработка POST запроса регистрации клиента"""
|
||||
form_data = await request.form()
|
||||
client_name = form_data.get("client_name")
|
||||
number_qubits = form_data.get("number_qubits")
|
||||
|
||||
if not number_qubits:
|
||||
return Div(
|
||||
P("❌ Обязательно необходимо указать колчество кубит", style="color: red;"),
|
||||
)
|
||||
|
||||
if not client_name:
|
||||
return Div(
|
||||
P("❌ Имя клиента обязательно", style="color: red;"),
|
||||
)
|
||||
|
||||
# Проверка формата имени клиента
|
||||
import re
|
||||
|
||||
if not re.match(r"^[a-zA-Z0-9_-]+$", client_name):
|
||||
return Div(
|
||||
P(
|
||||
"❌ Имя клиента может содержать только буквы, цифры, дефисы и подчеркивания",
|
||||
style="color: red;",
|
||||
),
|
||||
Button("Попробовать снова", onclick="location.reload()", cls="button"),
|
||||
)
|
||||
|
||||
token_data = load_token()
|
||||
if not token_data:
|
||||
return Div(
|
||||
P("❌ Сессия истекла. Пожалуйста, войдите снова.", style="color: red;"),
|
||||
A("Войти", href="/login", cls="button"),
|
||||
)
|
||||
|
||||
access_token = token_data.get("access_token")
|
||||
|
||||
# Регистрация в центральном микросервисе
|
||||
try:
|
||||
response = get_or_create_device_by_name(
|
||||
client_name, number_qubits, access_token
|
||||
)
|
||||
save_client_info(
|
||||
{
|
||||
"system_id": response["system_id"],
|
||||
}
|
||||
)
|
||||
|
||||
return Div(
|
||||
P(
|
||||
"✅ Клиент успешно зарегистрирован!",
|
||||
style="color: green; font-weight: bold;",
|
||||
),
|
||||
P(f"Имя клиента: {client_name}"),
|
||||
P("Перенаправление на панель управления..."),
|
||||
Meta(http_equiv="refresh", content="2;url=/dashboard"),
|
||||
Script("setTimeout(() => { window.location.href = '/dashboard'; }, 2000);"),
|
||||
)
|
||||
except Exception:
|
||||
return Div(
|
||||
P("❌ Ошибка регистрации", style="color: red;"),
|
||||
Button("Попробовать снова", onclick="location.reload()", cls="button"),
|
||||
P(
|
||||
"Если проблема повторяется, обратитесь к администратору.",
|
||||
style="font-size: 12px; margin-top: 10px;",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@app.route("/dashboard", methods="get")
|
||||
async def dashboard():
|
||||
access_token = get_valid_access_token()
|
||||
if rabbitmq_manager._connection and access_token:
|
||||
await rabbitmq_manager._connection.update_secret(
|
||||
access_token, reason="Token expired"
|
||||
)
|
||||
if not access_token:
|
||||
# Token refresh failed, redirect to login
|
||||
return RedirectResponse("/", status_code=303)
|
||||
token_data = load_token()
|
||||
client_info = load_client_info()
|
||||
|
||||
if not token_data or not client_info:
|
||||
return RedirectResponse("/", status_code=303)
|
||||
|
||||
user_info = token_data.get("user_info", {})
|
||||
username = user_info.get(
|
||||
"preferred_username", user_info.get("name", "Пользователь")
|
||||
)
|
||||
system_id = client_info.get("system_id", None)
|
||||
|
||||
if not token_data:
|
||||
return Div(
|
||||
P("❌ Сессия истекла. Пожалуйста, войдите снова.", style="color: red;"),
|
||||
A("Войти", href="/login", cls="button"),
|
||||
)
|
||||
try:
|
||||
if system_id:
|
||||
device_data = get_device_by_id(system_id, access_token)
|
||||
device_name = device_data["system"]["system_name"]
|
||||
max_qubits = device_data["system"]["max_qubits"]
|
||||
created_at = device_data["system"]["created_at"]
|
||||
else:
|
||||
device_name = "error"
|
||||
max_qubits = "error"
|
||||
created_at = "error"
|
||||
except Exception as a:
|
||||
device_name = "error"
|
||||
max_qubits = "error"
|
||||
created_at = "error"
|
||||
|
||||
return Titled(
|
||||
"Панель управления",
|
||||
Container(
|
||||
Card(
|
||||
P(Strong("Имя пользователя: "), username),
|
||||
P(Strong("Email: "), user_info.get("email", "Не указан")),
|
||||
P(Strong("Устройство зарегестрировано: "), created_at),
|
||||
Div(style="border: 1px solid black"),
|
||||
P(Strong("Имя вычислительной системы: "), device_name),
|
||||
P(Strong("Макс. количество кубит: "), max_qubits),
|
||||
Div(Button("Изменить", style="background-color: blue;")),
|
||||
Div(
|
||||
id="client-status",
|
||||
hx_get="/client-status",
|
||||
hx_trigger="load, every 30s",
|
||||
),
|
||||
),
|
||||
Div(id="memory-stats", hx_get="/memory-stats", hx_trigger="load, every 5s"),
|
||||
A(
|
||||
Button(
|
||||
"Выйти",
|
||||
style="background-color: #f44336;",
|
||||
),
|
||||
href="/logout",
|
||||
style="margin: 0px;color: inherit;text-decoration: inherit;",
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@app.route("/client-status", methods="get")
|
||||
async def status():
|
||||
"""Получить статус клиента из центрального микросервиса"""
|
||||
token_data = load_token()
|
||||
client_info = load_client_info()
|
||||
|
||||
if not token_data or not client_info:
|
||||
return P("Клиент не зарегистрирован", style="color: orange;")
|
||||
|
||||
if rabbitmq_manager._connection:
|
||||
# TODO: FIX
|
||||
status = (
|
||||
rabbitmq_manager._connection.connected
|
||||
) # get_client_status(client_name, access_token)
|
||||
|
||||
if status:
|
||||
return Div(
|
||||
P(
|
||||
Strong("Статус клиента: "),
|
||||
Span("✅ Активен", style="color: green;"),
|
||||
),
|
||||
style="padding: 10px; border-radius: 5px; margin: 10px 0;",
|
||||
)
|
||||
else:
|
||||
return Div(
|
||||
P(
|
||||
Strong("Статус клиента: "),
|
||||
Span("Не подключен", style="color: orange;"),
|
||||
),
|
||||
P(
|
||||
"Не удалось связаться с центральным сервисом или данному клиенту отказано в подключении"
|
||||
),
|
||||
style="padding: 10px; border-radius: 5px; margin: 10px 0;",
|
||||
)
|
||||
|
||||
|
||||
@app.route("/memory-stats", methods="get")
|
||||
async def memory_stats():
|
||||
"""Endpoint, возвращающий карточку использования памяти (для обновления через HTMX)"""
|
||||
process = psutil.Process(os.getpid())
|
||||
memory_info = process.memory_info()
|
||||
current_time = datetime.now()
|
||||
timestamp = current_time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
def format_bytes(bytes):
|
||||
for unit in ["Б", "КБ", "МБ", "ГБ"]:
|
||||
if bytes < 1024.0:
|
||||
return f"{bytes:.1f} {unit}"
|
||||
bytes /= 1024.0
|
||||
return f"{bytes:.1f} ГБ"
|
||||
|
||||
mem = 0
|
||||
if os.path.isfile("/sys/fs/cgroup/memory.max"):
|
||||
with open("/sys/fs/cgroup/memory.max") as limit:
|
||||
try:
|
||||
mem = int(limit.read())
|
||||
except Exception:
|
||||
mem = 0
|
||||
# resource.setrlimit(resource.RLIMIT_AS, (mem, mem))
|
||||
|
||||
rss_memory = format_bytes(memory_info.rss)
|
||||
if mem > 0:
|
||||
memory_percent = (memory_info.rss / mem) * 100
|
||||
else:
|
||||
memory_percent = process.memory_percent()
|
||||
|
||||
bar_color = (
|
||||
"#4CAF50"
|
||||
if memory_percent < 1
|
||||
else "#FFC107"
|
||||
if memory_percent < 5
|
||||
else "#F44336"
|
||||
)
|
||||
|
||||
return Div(
|
||||
H2("Использование памяти"),
|
||||
Card(
|
||||
Div(
|
||||
H3(rss_memory),
|
||||
style="width: 100%;display:flex;flex-direction:column;align-items:center",
|
||||
),
|
||||
Div(
|
||||
Div(
|
||||
style=f"width: {min(memory_percent, 100)}%; background-color: {bar_color}; height: 20px; border-radius: 10px;"
|
||||
),
|
||||
style="width: 100%; background-color: #e0e0e0; border-radius: 10px; overflow: hidden; margin: 10px 0;",
|
||||
),
|
||||
P(
|
||||
f"{memory_percent:.2f}% от системной памяти",
|
||||
style="text-align: center; font-size: 12px;",
|
||||
),
|
||||
P(
|
||||
Strong("Обновлено: "),
|
||||
timestamp,
|
||||
style="text-align: center; font-size: 11px; color: #666;",
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@app.route("/logout", methods="get")
|
||||
async def logout():
|
||||
token_data = load_token()
|
||||
|
||||
if token_data and token_data.get("refresh_token"):
|
||||
keycloak_openid.logout(token_data["refresh_token"])
|
||||
# Очищаем локальные файлы
|
||||
delete_token()
|
||||
delete_client_info()
|
||||
delete_verifier()
|
||||
|
||||
return Titled(
|
||||
"Выход выполнен",
|
||||
H2("До свидания!"),
|
||||
P("Вы успешно вышли из системы."),
|
||||
P("Ваш клиент был отменен в центральном сервисе."),
|
||||
A("Войти снова", href="/", cls="button"),
|
||||
)
|
||||
|
||||
|
||||
# Опционально: endpoint для проверки работоспособности для центрального сервиса
|
||||
@app.route("/health", methods="get")
|
||||
async def health():
|
||||
return {"status": "healthy", "timestamp": datetime.now().isoformat()}
|
||||
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()
|
||||
168
Dockerfile_build/source/modules/vqe.py
Normal file
168
Dockerfile_build/source/modules/vqe.py
Normal file
@@ -0,0 +1,168 @@
|
||||
import os
|
||||
from multiprocessing.connection import Connection
|
||||
|
||||
import jax
|
||||
import pennylane as qml
|
||||
import pennylane.numpy as np
|
||||
from jax import numpy as jnp
|
||||
from pennylane import qchem
|
||||
from pennylane.devices import Device
|
||||
from pennylane.optimize import GradientDescentOptimizer
|
||||
|
||||
jax.config.update("jax_enable_x64", True)
|
||||
|
||||
os.environ["OMP_NUM_THREADS"] = "16"
|
||||
|
||||
|
||||
def parse_xyz_from_text(text: str):
|
||||
"""Parse XYZ format from text content."""
|
||||
lines = text.strip().split("\n")
|
||||
|
||||
# First line: number of atoms
|
||||
num_atoms = int(lines[0].strip())
|
||||
|
||||
# Second line: Charge/Multiplicity/Electrons/Orbitals (optional)
|
||||
# Skip or parse as needed
|
||||
|
||||
symbols = []
|
||||
coordinates = []
|
||||
|
||||
# Parse atom lines (after the second line)
|
||||
for line in lines[2 : 2 + num_atoms]:
|
||||
parts = line.strip().split()
|
||||
if len(parts) >= 4:
|
||||
symbol = parts[0]
|
||||
x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
|
||||
symbols.append(symbol)
|
||||
coordinates.append([x, y, z])
|
||||
|
||||
return symbols, coordinates
|
||||
|
||||
|
||||
def extract_electron_info(text: str):
|
||||
"""Extract electron and orbital counts from the second line."""
|
||||
lines = text.strip().split("\n")
|
||||
if len(lines) >= 2:
|
||||
second_line = lines[1]
|
||||
# Parse "Charge=0 Multiplicity=1 Electrons=3 Orbitals=3"
|
||||
electrons = 3 # default
|
||||
orbitals = 3 # default
|
||||
charge = 0 # default
|
||||
multiplicity = 1 # default
|
||||
|
||||
for part in second_line.split():
|
||||
if "Electrons=" in part:
|
||||
electrons = int(part.split("=")[1])
|
||||
elif "Orbitals=" in part:
|
||||
orbitals = int(part.split("=")[1])
|
||||
elif "Charge=" in part:
|
||||
charge = int(part.split("=")[1])
|
||||
elif "Multiplicity=" in part:
|
||||
multiplicity = int(part.split("=")[1])
|
||||
|
||||
return electrons, orbitals, charge, multiplicity
|
||||
|
||||
return 3, 3, 0, 1 # fallback defaults
|
||||
|
||||
|
||||
def prepare_data(data):
|
||||
text_content = data.get("text", "")
|
||||
|
||||
# Parse the molecular data (assuming it's in XYZ format)
|
||||
symbols, coordinates = parse_xyz_from_text(text_content)
|
||||
|
||||
# Extract electron/orbital info (from the Charge/Multiplicity line)
|
||||
# "Charge=0 Multiplicity=1 Electrons=3 Orbitals=3"
|
||||
electrons, orbitals, charge, multiplicity = extract_electron_info(text_content)
|
||||
|
||||
return {
|
||||
"symbols": symbols,
|
||||
"coordinates": coordinates,
|
||||
"charge": charge,
|
||||
"multiplicity": multiplicity,
|
||||
"active_electrons": electrons,
|
||||
"active_orbitals": orbitals,
|
||||
"max_iterations": data.get("max_iterations", 200),
|
||||
"conv_tol": data.get("conv_tol", 1e-6),
|
||||
"step_size": data.get("step_size", 0.05),
|
||||
}
|
||||
|
||||
|
||||
def run_vqe(conn: Connection, dev1: Device, data: dict):
|
||||
coordinates = jnp.array(data.get("coordinates"))
|
||||
charge = int(data.get("charge"))
|
||||
multiplicity = int(data.get("multiplicity"))
|
||||
molecule = qchem.Molecule(
|
||||
data.get("symbols"),
|
||||
coordinates,
|
||||
charge=charge,
|
||||
mult=multiplicity,
|
||||
)
|
||||
|
||||
active_electrons = int(data.get("active_electrons"))
|
||||
active_orbitals = int(data.get("active_orbitals"))
|
||||
|
||||
max_iterations = int(data.get("max_iterations", 200))
|
||||
step_size = float(data.get("step_size", 0.05))
|
||||
conv_tol = float(data.get("conv_tol", 1e-6))
|
||||
|
||||
H, qubits = qchem.molecular_hamiltonian(
|
||||
molecule,
|
||||
active_electrons=active_electrons,
|
||||
active_orbitals=active_orbitals,
|
||||
method="openfermion",
|
||||
) # type: ignore
|
||||
|
||||
singles, doubles = qml.qchem.excitations(active_electrons, qubits)
|
||||
|
||||
params = np.array(np.zeros(len(singles) + len(doubles)), requires_grad=True)
|
||||
|
||||
conn.send(
|
||||
{
|
||||
"iter_num": 0,
|
||||
"energy": None,
|
||||
"conv": None,
|
||||
"params": params.tolist() if hasattr(params, "tolist") else list(params),
|
||||
}
|
||||
)
|
||||
|
||||
@qml.qnode(dev1)
|
||||
def circuit(param, wires):
|
||||
# Map excitations to the wires the UCCSD circuit will act on
|
||||
s_wires, d_wires = qml.qchem.excitations_to_wires(singles, doubles)
|
||||
qml.UCCSD(
|
||||
param,
|
||||
wires,
|
||||
s_wires=s_wires,
|
||||
d_wires=d_wires,
|
||||
init_state=qml.qchem.hf_state(active_electrons, qubits),
|
||||
)
|
||||
return qml.expval(H)
|
||||
|
||||
def cost_fn(param):
|
||||
return circuit(param, wires=range(qubits))
|
||||
|
||||
opt = GradientDescentOptimizer(stepsize=step_size)
|
||||
|
||||
for n in range(max_iterations):
|
||||
# Take step
|
||||
params, prev_energy = opt.step_and_cost(cost_fn, params)
|
||||
|
||||
energy = cost_fn(params)
|
||||
|
||||
# Calculate difference between new and old energies
|
||||
conv = np.abs(energy - prev_energy)
|
||||
|
||||
conn.send(
|
||||
{
|
||||
"iter_num": n,
|
||||
"energy": float(energy),
|
||||
"conv": float(conv),
|
||||
"params": params.tolist()
|
||||
if hasattr(params, "tolist")
|
||||
else list(params),
|
||||
}
|
||||
)
|
||||
|
||||
if conv <= conv_tol:
|
||||
break
|
||||
Reference in New Issue
Block a user