v1.0
All checks were successful
Build and Deploy Docker Image / build-and-push (push) Successful in 5m5s

working basic setup

- added profile picture upload
- fied a lot of errors with rabbitmq
- updated the "module" tables
- added experimental runner
This commit is contained in:
2026-05-24 22:15:30 +03:00
parent ea4d91f3bd
commit a571ef9c7a
19 changed files with 1776 additions and 329 deletions

View File

@@ -8,3 +8,6 @@ asyncpg==0.31.0
python-keycloak==7.1.1
pydantic==2.13.1
minio==7.2.20
aio-pika==9.5.8

View File

@@ -1,10 +1,20 @@
import json
from datetime import datetime
from typing import List
from typing import List, Optional
from connections.db import get_db
from connections.keycloak import (
get_current_token_payload,
)
from connections.minio import (
ensure_bucket_exists,
get_file_from_minio,
upload_file_to_minio,
)
from connections.rabbitmq import (
bind_queues_to_team_exchange,
publish_task_to_team,
)
from crud.experiment_crud import (
create_experiment,
create_experiment_type,
@@ -13,49 +23,141 @@ from crud.experiment_crud import (
delete_instance,
get_all_experiment_types,
get_experiment_instances,
get_experiment_type_by_id,
get_single_experiment,
get_single_instance,
get_user_experiments,
set_simulation_result,
update_experiment,
update_experiment_type,
update_instance,
)
from crud.team_crud import check_team_permission
from fastapi import Depends, Query
# Get team's systems to find max qubits available
from crud.team_crud import (
check_team_permission,
get_team, # You'll need to import this
)
from fastapi import Depends, File, Query, UploadFile
from fastapi.param_functions import Form
from fastapi.routing import APIRouter
from rest_models.experiment_models import (
CreateExperimentRequest,
CreateExperimentResponse,
CreateExperimentTypeRequest,
CreateExperimentTypeResponse,
CreateInstanceRequest,
CreateInstanceResponse,
ExperimentData,
ExperimentListResponse,
ExperimentTypeData,
ExperimentTypeList,
InstanceData,
InstanceListResponse,
SimpleInstanceData,
SimulationResultData,
StartExperimentRequest,
StartExperimentResponse,
UpdateExperimentRequest,
UpdateExperimentResponse,
UpdateInstanceRequest,
UpdateInstanceResponse,
)
from rest_models.machine_models import (
ComputationalSystemShortData,
)
from rest_models.team_models import (
TeamsShortListResponse,
)
from sqlalchemy.ext.asyncio import AsyncSession
from starlette.exceptions import HTTPException
from starlette.responses import StreamingResponse
from starlette.status import HTTP_200_OK
router = APIRouter(tags=["Experiments"])
# ============= HELPER FUNC ===========================
#
async def validate_instance_qubits(
db: AsyncSession, team_id: int, qubits_needed: int
) -> None:
"""
Validate that qubits_needed is within allowed range for the team.
Args:
db: Database session
team_id: Team ID to check systems for
qubits_needed: Number of qubits needed for the instance
Raises:
HTTPException: If validation fails
"""
# Check qubits is positive
if qubits_needed <= 0:
raise HTTPException(
400,
f"Invalid qubits_needed: {qubits_needed}. Qubits must be greater than 0.",
)
team_systems = await get_team(db, team_id)
if not team_systems:
raise HTTPException(
400,
"Team has no computational systems configured. "
"Cannot validate qubits requirement.",
)
max_qubits = max(system.qubits_given for system in team_systems.team_systems)
if qubits_needed > max_qubits:
raise HTTPException(
400,
f"Instance requires {qubits_needed} qubits, but the maximum available "
f"in team's systems is {max_qubits} qubits.",
)
def get_experiment_collective_status(simulation_statuses: list[str]) -> str:
"""
Get the collective status of an experiment based on its simulation statuses.
Args:
simulation_statuses: List of simulation status strings (e.g., ["DRAFT", "COMPLETE", ...])
Returns:
Collective experiment status: "DRAFT", "PROCESSING", "IN QUEUE", or "COMPLETE"
"""
if not simulation_statuses:
return "DRAFT"
# Check if all are DRAFT (or not set/empty)
if all(status == "DRAFT" for status in simulation_statuses):
return "DRAFT"
# Check if any is PROCESSING
if any(
status == "PROCESSING" or status == "IN SYSTEM"
for status in simulation_statuses
):
return "RUNNING"
# Check if any is IN QUEUE (and none are PROCESSING)
if any(status == "IN QUEUE" for status in simulation_statuses):
return "IN QUEUE"
# If all are COMPLETE or ERROR
if all(status in ["COMPLETE", "ERROR"] for status in simulation_statuses):
if any(status == "ERROR" for status in simulation_statuses):
return "COMPLETE WITH ERROR"
return "COMPLETE"
# Default fallback (should not happen with valid statuses)
return "DRAFT"
# ============= EXPERIMENT TYPE ENDPOINTS =============
@router.get("/types", response_model=List[ExperimentTypeData])
@router.get("/types", response_model=List[ExperimentTypeList])
async def get_experiment_types_request(
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> List[ExperimentTypeData]:
) -> List[ExperimentTypeList]:
"""Get all existing experiment types"""
keycloak_id = payload.get("sub")
@@ -65,13 +167,10 @@ async def get_experiment_types_request(
experiment_types = await get_all_experiment_types(db)
return [
ExperimentTypeData(
ExperimentTypeList(
id=exp_type.id,
name=exp_type.name,
description=exp_type.description,
file_frontend=exp_type.file_frontend,
server_path=exp_type.server_path,
file_comp_system=exp_type.file_comp_system,
)
for exp_type in experiment_types
]
@@ -79,46 +178,184 @@ async def get_experiment_types_request(
@router.post("/types", response_model=CreateExperimentTypeResponse)
async def create_experiment_type_request(
create_data: CreateExperimentTypeRequest,
name: str = Form(),
file_frontend: UploadFile = File(...),
file_comp_system: UploadFile = File(...),
description: Optional[str] = None,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> CreateExperimentTypeResponse:
"""Create a new experiment type (admin only)"""
keycloak_id = payload.get("sub")
"""Create a new experiment type with file uploads to MinIO (admin only)"""
# Check authentication
keycloak_id = payload.get("sub")
if not keycloak_id:
raise HTTPException(403, "Permission denied")
raise HTTPException(status_code=403, detail="Permission denied")
# TODO: Add admin permission check here if needed
# Ensure MinIO bucket exists
await ensure_bucket_exists()
# First create the experiment type in DB to get ID
experiment_type = await create_experiment_type(
db=db,
name=create_data.name,
file_frontend=create_data.file_frontend,
server_path=create_data.server_path,
file_comp_system=create_data.file_comp_system,
description=create_data.description,
name=name,
description=description,
file_frontend="", # Temporary, will update
file_comp_system="", # Temporary, will update
)
try:
# Upload files using the new ID
frontend_path = await upload_file_to_minio(
file=file_frontend, experiment_type_id=experiment_type.id
)
comp_system_path = await upload_file_to_minio(
file=file_comp_system, experiment_type_id=experiment_type.id
)
# Update experiment type with file paths
experiment_type.file_frontend = frontend_path
experiment_type.file_comp_system = comp_system_path
await db.commit()
await db.refresh(experiment_type)
return CreateExperimentTypeResponse(
id=experiment_type.id,
name=experiment_type.name,
description=experiment_type.description,
)
except Exception as e:
# Rollback if upload fails
await db.rollback()
raise HTTPException(
500,
detail=f"Failed to create experiment type: {str(e)}",
)
@router.put("/types", response_model=CreateExperimentTypeResponse)
async def update_type_request(
type_id: int = Form(),
name: Optional[str] = Form(None),
file_frontend: Optional[UploadFile] | None = File(None),
file_comp_system: Optional[UploadFile] | None = File(None),
description: Optional[str] = None,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> CreateExperimentTypeResponse:
"""Create a new experiment type with file uploads to MinIO (admin only)"""
# Check authentication
keycloak_id = payload.get("sub")
if not keycloak_id:
raise HTTPException(status_code=403, detail="Permission denied")
# Ensure MinIO bucket exists
await ensure_bucket_exists()
frontend_path = None
comp_path = None
if file_frontend:
frontend_path = await upload_file_to_minio(
file=file_frontend, experiment_type_id=type_id
)
if file_comp_system:
comp_path = await upload_file_to_minio(
file=file_comp_system, experiment_type_id=type_id
)
# First create the experiment type in DB to get ID
experiment_type = await update_experiment_type(
db=db,
id=type_id,
name=name,
description=description,
file_frontend=frontend_path, # Temporary, will update
file_comp_system=comp_path, # Temporary, will update
)
return CreateExperimentTypeResponse(
id=experiment_type.id,
name=experiment_type.name,
description=experiment_type.description,
file_frontend=experiment_type.file_frontend,
server_path=experiment_type.server_path,
file_comp_system=experiment_type.file_comp_system,
)
@router.get("/types/frontend")
async def get_frontend_file(
experiment_type_id: int = Query(),
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
):
"""Get the frontend HTML file"""
# Get experiment type from database
experiment_type = await get_experiment_type_by_id(db, experiment_type_id)
if not experiment_type:
raise HTTPException(status_code=404, detail="Experiment type not found")
if not experiment_type.file_frontend:
raise HTTPException(status_code=404, detail="Frontend file not found")
# Get file from MinIO
file_response = await get_file_from_minio(experiment_type.file_frontend)
# Return file as streaming response with correct content type
return StreamingResponse(
file_response,
media_type="text/html",
headers={
"Content-Disposition": f"inline; filename=frontend_{experiment_type_id}.html"
},
)
@router.get("/types/comp-system")
async def get_comp_system_file(
experiment_type_id: int,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
):
"""Get the comp system Python file"""
# Get experiment type from database
experiment_type = await get_experiment_type_by_id(db, experiment_type_id)
if not experiment_type:
raise HTTPException(status_code=404, detail="Experiment type not found")
if not experiment_type.file_comp_system:
raise HTTPException(status_code=404, detail="Comp system file not found")
# Get file from MinIO
file_response = await get_file_from_minio(experiment_type.file_comp_system)
# Return Python file for download (since it will be imported dynamically)
return StreamingResponse(
file_response,
media_type="text/x-python",
headers={
"Content-Disposition": f"attachment; filename=comp_system_{experiment_type_id}.py"
},
)
# ============= EXPERIMENT ENDPOINTS =============
@router.post("", response_model=CreateExperimentResponse)
@router.post("", response_model=ExperimentData)
async def create_experiment_request(
create_data: CreateExperimentRequest,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> CreateExperimentResponse:
) -> ExperimentData:
"""Create a new experiment"""
keycloak_id = payload.get("sub")
@@ -138,13 +375,44 @@ async def create_experiment_request(
description=create_data.description,
)
return CreateExperimentResponse(
return ExperimentData(
id=experiment.id,
team_id=experiment.team_id,
experiment_type_id=experiment.experiment_type_id,
team=TeamsShortListResponse(
team_id=experiment.team.id, team_name=experiment.team.name
),
name=experiment.name,
description=experiment.description,
created_at=experiment.created_at,
experiment_type=ExperimentTypeList(
id=experiment.experiment_type.id,
name=experiment.experiment_type.name,
description=experiment.experiment_type.description,
),
instances_count=0,
status=get_experiment_collective_status(
list(
map(
lambda x: (
x.simulation_result.status.name
if x.simulation_result
else "DRAFT"
),
experiment.instances,
)
)
),
instance_preview=list(
map(
lambda inst: SimpleInstanceData(
id=inst.id,
instance_data={},
name=inst.name,
qubits_needed=inst.qubits_needed,
description=inst.description,
),
experiment.instances[:6],
)
),
)
@@ -169,25 +437,51 @@ async def get_experiment_request(
return ExperimentData(
id=experiment.id,
team_id=experiment.team_id,
team_name=experiment.team.name if experiment.team else None,
experiment_type_id=experiment.experiment_type_id,
experiment_type_name=experiment.experiment_type.name
if experiment.experiment_type
else None,
team=TeamsShortListResponse(
team_id=experiment.team.id, team_name=experiment.team.name
),
name=experiment.name,
description=experiment.description,
created_at=experiment.created_at,
instances_count=len(experiment.instances) if experiment.instances else 0,
experiment_type=ExperimentTypeList(
id=experiment.experiment_type.id,
name=experiment.experiment_type.name,
description=experiment.experiment_type.description,
),
instances_count=len(experiment.instances),
status=get_experiment_collective_status(
list(
map(
lambda x: (
x.simulation_result.status.name
if x.simulation_result
else "DRAFT"
),
experiment.instances,
)
)
),
instance_preview=list(
map(
lambda inst: SimpleInstanceData(
id=inst.id,
instance_data="{}",
name=inst.name,
qubits_needed=inst.qubits_needed,
description=inst.description,
),
experiment.instances[:6],
)
),
)
@router.put("", response_model=UpdateExperimentResponse)
@router.put("", response_model=ExperimentData)
async def update_experiment_request(
update_data: UpdateExperimentRequest,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> UpdateExperimentResponse:
) -> ExperimentData:
"""Update experiment data (cannot update type)"""
keycloak_id = payload.get("sub")
@@ -210,13 +504,44 @@ async def update_experiment_request(
description=update_data.description,
)
return UpdateExperimentResponse(
return ExperimentData(
id=updated_experiment.id,
team_id=updated_experiment.team_id,
experiment_type_id=updated_experiment.experiment_type_id,
team=TeamsShortListResponse(
team_id=updated_experiment.team.id, team_name=updated_experiment.team.name
),
name=updated_experiment.name,
description=updated_experiment.description,
created_at=updated_experiment.created_at,
experiment_type=ExperimentTypeList(
id=updated_experiment.experiment_type.id,
name=updated_experiment.experiment_type.name,
description=updated_experiment.experiment_type.description,
),
instances_count=len(experiment.instances),
status=get_experiment_collective_status(
list(
map(
lambda x: (
x.simulation_result.status.name
if x.simulation_result
else "DRAFT"
),
experiment.instances,
)
)
),
instance_preview=list(
map(
lambda inst: SimpleInstanceData(
id=inst.id,
instance_data=json.dumps(inst.instance_data),
name=inst.name,
qubits_needed=inst.qubits_needed,
description=inst.description,
),
updated_experiment.instances[:6],
)
),
)
@@ -244,18 +569,42 @@ async def get_user_experiments_request(
result.append(
ExperimentData(
id=experiment.id,
team_id=experiment.team_id,
team_name=experiment.team.name if experiment.team else None,
experiment_type_id=experiment.experiment_type_id,
experiment_type_name=experiment.experiment_type.name
if experiment.experiment_type
else None,
team=TeamsShortListResponse(
team_id=experiment.team.id, team_name=experiment.team.name
),
name=experiment.name,
description=experiment.description,
created_at=experiment.created_at,
instances_count=len(experiment.instances)
if experiment.instances
else 0,
experiment_type=ExperimentTypeList(
id=experiment.experiment_type.id,
name=experiment.experiment_type.name,
description=experiment.experiment_type.description,
),
instances_count=len(experiment.instances),
status=get_experiment_collective_status(
list(
map(
lambda x: (
x.simulation_result.status.name
if x.simulation_result
else "DRAFT"
),
experiment.instances,
)
)
),
instance_preview=list(
map(
lambda inst: SimpleInstanceData(
id=inst.id,
instance_data="{}",
name=inst.name,
qubits_needed=inst.qubits_needed,
description=inst.description,
),
experiment.instances[:6],
)
),
)
)
@@ -299,12 +648,12 @@ async def delete_experiment_request(
# ============= INSTANCE ENDPOINTS =============
@router.post("/instance", response_model=CreateInstanceResponse)
@router.post("/instance", response_model=SimpleInstanceData)
async def create_instance_request(
create_data: CreateInstanceRequest,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> CreateInstanceResponse:
) -> SimpleInstanceData:
"""Add an instance to an experiment"""
keycloak_id = payload.get("sub")
@@ -323,27 +672,26 @@ async def create_instance_request(
instance = await create_instance(
db=db,
experiment_id=create_data.experiment_id,
instance_data_id=create_data.instance_data_id,
instance_data=create_data.instance_data,
name=create_data.name,
description=create_data.description,
)
return CreateInstanceResponse(
return SimpleInstanceData(
id=instance.id,
experiment_id=instance.experiment_id,
instance_data_id=instance.instance_data_id,
instance_data=json.dumps(create_data.instance_data),
name=instance.name,
qubits_needed=instance.qubits_needed,
description=instance.description,
simulation_result_id=instance.simulation_result_id,
)
@router.put("/instance", response_model=UpdateInstanceResponse)
@router.put("/instance", response_model=SimpleInstanceData)
async def update_instance_request(
update_data: UpdateInstanceRequest,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> UpdateInstanceResponse:
) -> SimpleInstanceData:
"""Update instance data (cannot update instance_data_id)"""
keycloak_id = payload.get("sub")
@@ -364,20 +712,26 @@ async def update_instance_request(
["manage_experiments"], experiment.team_id, db, keycloak_id
)
if update_data.qubits_needed is not None:
await validate_instance_qubits(
db, experiment.team_id, update_data.qubits_needed
)
updated_instance = await update_instance(
db=db,
instance_id=update_data.instance_id,
name=update_data.name,
description=update_data.description,
instance_data=update_data.instance_data,
qubits_needed=update_data.qubits_needed,
)
return UpdateInstanceResponse(
return SimpleInstanceData(
id=updated_instance.id,
experiment_id=updated_instance.experiment_id,
instance_data_id=updated_instance.instance_data_id,
instance_data=json.dumps(updated_instance.instance_data),
name=updated_instance.name,
qubits_needed=updated_instance.qubits_needed,
description=updated_instance.description,
simulation_result_id=updated_instance.simulation_result_id,
)
@@ -412,19 +766,26 @@ async def get_experiment_instances_request(
for instance in instances:
result.append(
InstanceData(
id=instance.id,
experiment_id=instance.experiment_id,
instance_data_id=instance.instance_data_id,
instance_id=instance.id,
instance_data=json.dumps(instance.instance_data),
name=instance.name,
qubits_needed=instance.qubits_needed,
description=instance.description,
simulation_result_id=instance.simulation_result_id,
simulation_status=instance.simulation_result.status.name
if instance.simulation_result and instance.simulation_result.status
else None,
simulation_started_at=instance.simulation_result.started_at
if instance.simulation_result
else None,
simulation_ended_at=instance.simulation_result.ended_at
simulation_result=SimulationResultData(
id=instance.simulation_result.id,
comp_system=ComputationalSystemShortData(
system_id=instance.simulation_result.computational_system.id,
system_name=instance.simulation_result.computational_system.system_name,
)
if instance.simulation_result.computational_system
else None,
simulation_result=json.dumps(
instance.simulation_result.simulation_result
),
status=instance.simulation_result.status.name,
started_at=instance.simulation_result.started_at,
ended_at=instance.simulation_result.ended_at,
)
if instance.simulation_result
else None,
)
@@ -472,61 +833,7 @@ async def delete_instance_request(
return {"message": "Instance deleted successfully"}
# ============= SIMULATION ENDPOINTS =============
@router.post("/start", response_model=StartExperimentResponse)
async def start_experiment_request(
start_data: StartExperimentRequest,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
) -> StartExperimentResponse:
"""Start an experiment (create simulation result for an instance)"""
keycloak_id = payload.get("sub")
if not keycloak_id:
raise HTTPException(403, "Permission denied")
# Get instance to find its experiment
instance = await get_single_instance(db, start_data.instance_id)
if not instance:
raise HTTPException(404, "Instance not found")
experiment = await get_single_experiment(db, instance.experiment_id)
if not experiment:
raise HTTPException(404, "Experiment not found")
# Check if user has permission to run experiments in this team
await check_team_permission(
["run_experiments"], experiment.team_id, db, keycloak_id
)
# Create or update simulation result
simulation_result = await set_simulation_result(
db=db,
comp_system_id=start_data.comp_system_id,
simulation_result_id=start_data.simulation_result_id,
status_name="PENDING", # Start with PENDING status
started_at=datetime.now(),
)
# Link simulation result to instance if not already linked
if not instance.simulation_result_id:
await update_instance(
db=db,
instance_id=instance.id,
simulation_result_id=simulation_result.id,
)
return StartExperimentResponse(
instance_id=instance.id,
simulation_result_id=simulation_result.id,
status="PENDING",
started_at=simulation_result.started_at,
)
@router.get("/instance", response_model=InstanceData)
@router.get("/instance/id", response_model=InstanceData)
async def get_instance_request(
instance_id: int = Query(),
db: AsyncSession = Depends(get_db),
@@ -551,19 +858,111 @@ async def get_instance_request(
await check_team_permission([], experiment.team_id, db, keycloak_id)
return InstanceData(
id=instance.id,
experiment_id=instance.experiment_id,
instance_data_id=instance.instance_data_id,
instance_id=instance.id,
instance_data=json.dumps(instance.instance_data),
name=instance.name,
description=instance.description,
simulation_result_id=instance.simulation_result_id,
simulation_status=instance.simulation_result.status.name
if instance.simulation_result and instance.simulation_result.status
else None,
simulation_started_at=instance.simulation_result.started_at
if instance.simulation_result
else None,
simulation_ended_at=instance.simulation_result.ended_at
qubits_needed=instance.qubits_needed,
simulation_result=SimulationResultData(
id=instance.simulation_result.id,
comp_system=ComputationalSystemShortData(
system_id=instance.simulation_result.computational_system.id,
system_name=instance.simulation_result.computational_system.system_name,
)
if instance.simulation_result.computational_system
else None,
simulation_result=json.dumps(instance.simulation_result.simulation_result),
status=instance.simulation_result.status.name,
started_at=instance.simulation_result.started_at,
ended_at=instance.simulation_result.ended_at,
)
if instance.simulation_result
else None,
)
# ============= SIMULATION ENDPOINTS =============
@router.post("/start")
async def start_experiment_request(
start_data: StartExperimentRequest,
db: AsyncSession = Depends(get_db),
payload: dict = Depends(get_current_token_payload),
):
"""Start an experiment (create simulation result for an instance)"""
keycloak_id = payload.get("sub")
if not keycloak_id:
raise HTTPException(403, "Permission denied")
# Get instance to find its experiment
experiment = await get_single_experiment(db, start_data.experiment_id)
if not experiment:
raise HTTPException(404, "Experiment not found")
if not experiment.instances:
raise HTTPException(400, "No instances")
for task in experiment.instances:
if task.simulation_result:
# If simulation result exists, check its status
current_status = task.simulation_result.status.name
if current_status != "DRAFT":
raise HTTPException(
400,
f"Cannot start experiment. Task {task.id} has status '{current_status}'. "
f"Only tasks with 'draft' status can be started.",
)
# Check if user has permission to run experiments in this team
await check_team_permission(
["run_experiments"], experiment.team_id, db, keycloak_id
)
# TODO: add the starting rabbitMQ logic
max_qubits = max(i.qubits_given for i in experiment.team.team_systems)
for task in experiment.instances:
if task.qubits_needed > max_qubits or task.qubits_needed <= 0:
raise HTTPException(
400,
f"Cannot start experiment. Instance {task.id} requires {task.qubits_needed} qubits, "
f"but the maximum available in team's systems is {max_qubits} qubits.",
)
await bind_queues_to_team_exchange(experiment.team_id, max_qubits)
for task in sorted(
experiment.instances, key=lambda x: x.qubits_needed, reverse=True
):
try:
# Publish task to queue
await publish_task_to_team(task)
# Create simulation result with IN_QUEUE status
# Pass None for comp_system_id since it's not assigned yet
simulation_result = await set_simulation_result(
db=db,
comp_system_id=None, # No system assigned yet
instance_id=task.id,
status_name="IN QUEUE",
started_at=datetime.now(),
ended_at=None,
simulation_result_data={},
)
except Exception as e:
# Create failed simulation result
await set_simulation_result(
db=db,
comp_system_id=None,
instance_id=task.id,
status_name="ERROR",
started_at=datetime.now(),
ended_at=datetime.now(),
simulation_result_data={"error": str(e)},
)
raise HTTPException(500, f"Failed to publish task {task.id}: {str(e)}")
return HTTP_200_OK

View File

@@ -1,3 +1,5 @@
from typing import List
from connections.db import get_db
from connections.keycloak import (
KeycloakAdminService,
@@ -363,6 +365,30 @@ async def get_computational_system_request(
)
@router.get("/system/team", response_model=List[GiveSystemToTeamRequest])
async def get_computational_system_teams(
system_id: int = Query(),
db: AsyncSession = Depends(get_db),
) -> List[GiveSystemToTeamRequest]:
"""Get computational system by ID with all teams that have access (requires ownership or team access)"""
system = await get_computational_system(db, system_id)
if not system:
raise HTTPException(404, "Computational system not found")
teams_list = []
for team_system in system.team_systems:
teams_list.append(
GiveSystemToTeamRequest(
system_id=team_system.system_id,
team_id=team_system.team_id,
qubits_given=team_system.qubits_given,
)
)
return teams_list
@router.get("/status", response_model=SystemStatusResponse)
async def get_system_status_request(
system_id: int = Query(),

View File

@@ -4,8 +4,14 @@ from connections.keycloak import (
get_current_user,
get_keycloak_admin,
)
from connections.minio import (
delete_profile_picture,
ensure_bucket_exists,
get_profile_picture,
upload_profile_picture,
)
from crud.user_crud import get_or_create_user, update_user_profile
from fastapi import Depends
from fastapi import Depends, File, Response, UploadFile
from fastapi.exceptions import HTTPException
from fastapi.routing import APIRouter
from rest_models.user_models import (
@@ -13,6 +19,7 @@ from rest_models.user_models import (
UserUpdateRequest,
)
from sql_models.models import User
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
router = APIRouter()
@@ -95,3 +102,61 @@ async def get_user_by_email(
profile_picture_path=pfp, # Would need separate DB lookup
created_at=created_at,
)
@router.on_event("startup")
async def startup_event():
"""Ensure buckets exist on startup"""
await ensure_bucket_exists()
@router.post("/upload")
async def upload_profile_picture_endpoint(
file: UploadFile = File(...),
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user),
):
"""Upload a profile picture"""
# Upload to MinIO
object_path = await upload_profile_picture(
user_id=current_user.keycloak_id, file=file
)
# Update user record with the path
updated_user = await update_user_profile(db, current_user.keycloak_id, object_path)
if updated_user:
return {
"message": "Profile picture uploaded successfully",
"profile_picture_path": object_path,
}
else:
raise HTTPException(404, "Error getting user data")
# Optional: Public endpoint to view any user's profile picture
@router.get("/serve/{keycloak_id}")
async def get_user_profile_picture_endpoint(
keycloak_id: str,
db: AsyncSession = Depends(get_db),
current_user: User = Depends(
get_current_user
), # Require auth to prevent enumeration
):
"""Serve any user's profile picture (requires authentication)"""
result = await db.execute(select(User).where(User.keycloak_id == keycloak_id))
user = result.scalar_one_or_none()
if not user or not user.profile_picture_path:
# Return default profile picture or 404
raise HTTPException(404, "Profile picture not found")
# Get the file from MinIO
file_response = await get_profile_picture(user.profile_picture_path)
content = file_response.read()
return Response(
content=content,
media_type=file_response.headers.get("Content-Type", "image/jpeg"),
headers={"Cache-Control": "public, max-age=3600"},
)

View File

@@ -6,8 +6,18 @@ from api_endpoint.health_api import router as health_router
from api_endpoint.machine_api import router as machine_router
from api_endpoint.teams_api import router as team_router
from api_endpoint.user_api import router as user_router
from config.seeding import seed_permissions, seed_system_statuses
from config.logging_config import logging
from config.rabbitmq_config import rabbitmq_manager
from config.seeding import (
seed_permissions,
seed_simulation_statuses,
seed_system_statuses,
)
from connections.db import create_tables, engine
from connections.rabbitmq import (
start_heartbeat_monitoring,
start_progress_monitoring,
)
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
@@ -27,6 +37,13 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
async with AsyncSession(engine) as session:
await seed_permissions(session)
await seed_system_statuses(session)
await seed_simulation_statuses(session)
try:
await rabbitmq_manager.connect()
await start_heartbeat_monitoring(redis, session)
await start_progress_monitoring(session)
except Exception:
logging.error("Failed to connect to RabbitMQ")
yield
await engine.dispose()
@@ -35,10 +52,7 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
app = FastAPI(lifespan=lifespan)
origins = [
"http://localhost",
"http://localhost:8001",
]
origins = ["https://quantum.deowl.ru"]
app.add_middleware(
CORSMiddleware,

View File

@@ -0,0 +1,9 @@
# MinIO configuration
from minio.credentials.providers import os
MINIO_ENDPOINT = "minio:9000"
MINIO_ACCESS_KEY = os.getenv("MINIO_ACCESS_KEY", "minioadmin")
MINIO_SECRET_KEY = os.getenv("MINIO_SECRET_KEY", "minioadmin")
MINIO_SECURE = False
MINIO_BUCKET = "experiment-types"
USER_PFP_BUCKET = "user-pfp"

View File

@@ -0,0 +1,86 @@
import os
from typing import Optional
import aio_pika
from aio_pika.abc import (
AbstractChannel,
AbstractRobustConnection,
)
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):
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=os.environ["RABBITMQ_USER"],
password=os.environ["RABBITMQ_PASSWORD"],
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
else:
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 managing."""
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()

View File

@@ -1,5 +1,5 @@
from config.logging_config import logger
from sql_models.models import Permission, SystemStatus
from sql_models.models import Permission, SimulationStatus, SystemStatus
from sqlalchemy.ext.asyncio.session import AsyncSession
from sqlalchemy.sql import select
@@ -35,6 +35,27 @@ DEFAULT_STATUSES = [
{"name": "BUSY", "description": "System is busy processing other tasks"},
]
DEFAULT_SIMULATION_STATUSES = [
{
"name": "DRAFT",
"description": "Simulation is in draft state, not yet submitted for processing",
},
{
"name": "IN QUEUE",
"description": "Simulation is queued and waiting to be processed",
},
{
"name": "IN SYSTEM",
"description": "Simulation data has been picked up by system",
},
{"name": "PROCESSING", "description": "Simulation is currently being processed"},
{"name": "COMPLETE", "description": "Simulation has completed successfully"},
{
"name": "ERROR",
"description": "Simulation encountered an error during processing",
},
]
async def seed_permissions(db: AsyncSession) -> bool:
"""Seed default permissions into the database."""
@@ -78,3 +99,25 @@ async def seed_system_statuses(db: AsyncSession) -> bool:
logger.error(f"Error seeding system statuses: {e}")
await db.rollback()
raise
async def seed_simulation_statuses(db: AsyncSession) -> bool:
"""Seed default simulation statuses into the database."""
try:
result = await db.execute(select(SimulationStatus).limit(1))
if result.scalar_one_or_none():
logger.info("Simulation statuses already seeded, skipping...")
return False
statuses = [
SimulationStatus(name=status["name"], description=status["description"])
for status in DEFAULT_SIMULATION_STATUSES
]
db.add_all(statuses)
await db.commit()
logger.info(f"Seeded {len(statuses)} simulation statuses")
return True
except Exception as e:
logger.error(f"Error seeding simulation statuses: {e}")
await db.rollback()
raise

View File

@@ -0,0 +1,278 @@
# Create MinIO client
import io
from pathlib import Path
from config.minio_config import (
MINIO_ACCESS_KEY,
MINIO_BUCKET,
MINIO_ENDPOINT,
MINIO_SECRET_KEY,
MINIO_SECURE,
USER_PFP_BUCKET,
)
from fastapi import HTTPException, UploadFile
from minio import Minio
from minio.error import S3Error
minio_client = Minio(
MINIO_ENDPOINT,
access_key=MINIO_ACCESS_KEY,
secret_key=MINIO_SECRET_KEY,
secure=MINIO_SECURE,
)
async def ensure_bucket_exists():
"""Ensure that the bucket exists, create if it doesn't"""
try:
if not minio_client.bucket_exists(MINIO_BUCKET):
minio_client.make_bucket(MINIO_BUCKET)
print(f"Bucket '{MINIO_BUCKET}' created successfully")
else:
print(f"Bucket '{MINIO_BUCKET}' already exists")
except S3Error as err:
print(f"Error creating bucket: {err}")
raise
def detect_file_type(filename: str) -> str:
"""Detect file type from extension"""
extension = Path(filename).suffix.lower()
if extension in [".html", ".htm"]:
return "html"
elif extension in [".py"]:
return "python"
else:
return "unknown"
def get_content_type(filename: str) -> str:
"""Get appropriate content type for file"""
extension = Path(filename).suffix.lower()
if extension in [".html", ".htm"]:
return "text/html"
elif extension in [".py"]:
return "text/x-python"
else:
return "application/octet-stream"
async def upload_file_to_minio(
file: UploadFile,
experiment_type_id: int,
) -> str:
"""Upload a file to MinIO and return the object path"""
if not file.filename:
raise HTTPException(400, f"wrong file upload")
# Detect actual file format
detected_type = detect_file_type(file.filename)
# Generate object name using experiment_type_id
extension = Path(file.filename).suffix
object_name = f"experiment_types/{experiment_type_id}/{detected_type}{extension}"
# Read file content
content = await file.read()
file_size = len(content)
# Convert bytes to BytesIO (BinaryIO)
file_data = io.BytesIO(content)
try:
minio_client.remove_object(
bucket_name=MINIO_BUCKET,
object_name=object_name,
)
# Upload to MinIO
minio_client.put_object(
bucket_name=MINIO_BUCKET,
object_name=object_name,
data=file_data,
length=file_size,
content_type=get_content_type(file.filename),
)
return object_name
except Exception as e:
raise HTTPException(500, f"Failed to upload file to MinIO: {str(e)}")
async def get_file_from_minio(object_path: str):
"""Get file data from MinIO"""
try:
response = minio_client.get_object(
bucket_name=MINIO_BUCKET, object_name=object_path
)
return response
except Exception as e:
raise HTTPException(404, f"File not found: {str(e)}")
async def delete_file_from_minio(object_path: str) -> bool:
"""Delete a file from MinIO"""
try:
minio_client.remove_object(bucket_name=MINIO_BUCKET, object_name=object_path)
return True
except Exception as e:
print(f"Failed to delete file: {str(e)}")
return False
# ================ PFP =================
async def ensure_bucket_exists():
"""Ensure that the buckets exist, create if they don't"""
buckets = [MINIO_BUCKET, USER_PFP_BUCKET]
for bucket in buckets:
try:
if not minio_client.bucket_exists(bucket):
minio_client.make_bucket(bucket)
# Set bucket policy for public read (optional)
policy = {
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Principal": {"AWS": ["*"]},
"Action": ["s3:GetObject"],
"Resource": [f"arn:aws:s3:::{bucket}/*"],
}
],
}
# Uncomment if you want public read access
# minio_client.set_bucket_policy(bucket, json.dumps(policy))
print(f"Bucket '{bucket}' created successfully")
else:
print(f"Bucket '{bucket}' already exists")
except S3Error as err:
print(f"Error creating bucket {bucket}: {err}")
raise
def get_content_type_for_image(filename: str) -> str:
"""Get appropriate content type for image files"""
extension = Path(filename).suffix.lower()
content_types = {
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".png": "image/png",
".gif": "image/gif",
".webp": "image/webp",
".bmp": "image/bmp",
}
return content_types.get(extension, "application/octet-stream")
def validate_image_file(filename: str) -> bool:
"""Validate if the file is an allowed image type"""
allowed_extensions = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
extension = Path(filename).suffix.lower()
return extension in allowed_extensions
async def upload_profile_picture(
user_id: str, # Keycloak ID
file: UploadFile,
) -> str:
"""
Upload a user profile picture to MinIO and return the object path
"""
if not file.filename:
raise HTTPException(400, "No filename provided")
# Validate file type
if not validate_image_file(file.filename):
raise HTTPException(
400, "Invalid file type. Allowed types: jpg, jpeg, png, gif, webp, bmp"
)
# Validate file size (e.g., max 5MB)
file_size = 0
content = await file.read()
file_size = len(content)
max_size = 5 * 1024 * 1024 # 5MB
if file_size > max_size:
raise HTTPException(
400, f"File too large. Max size: {max_size // (1024 * 1024)}MB"
)
# Generate unique object name
extension = Path(file.filename).suffix.lower()
object_name = f"users/{user_id}/profile_picture{extension}"
# Convert bytes to BytesIO
file_data = io.BytesIO(content)
try:
# Delete existing profile picture if exists
try:
minio_client.remove_object(
bucket_name=USER_PFP_BUCKET,
object_name=object_name,
)
except S3Error:
# Object might not exist, continue
pass
# Upload to MinIO
minio_client.put_object(
bucket_name=USER_PFP_BUCKET,
object_name=object_name,
data=file_data,
length=file_size,
content_type=get_content_type_for_image(file.filename),
)
return object_name
except Exception as e:
raise HTTPException(500, f"Failed to upload profile picture: {str(e)}")
async def get_profile_picture_url(object_path: str) -> str:
"""
Generate a presigned URL for temporary access to the profile picture
"""
try:
# Generate URL that expires in 1 hour (3600 seconds)
url = minio_client.presigned_get_object(
bucket_name=USER_PFP_BUCKET,
object_name=object_path,
expires=3600, # 1 hour
)
return url
except S3Error as e:
raise HTTPException(404, f"Profile picture not found: {str(e)}")
async def get_profile_picture(object_path: str):
"""
Get the actual profile picture file data
"""
try:
response = minio_client.get_object(
bucket_name=USER_PFP_BUCKET, object_name=object_path
)
return response
except S3Error as e:
raise HTTPException(404, f"Profile picture not found: {str(e)}")
async def delete_profile_picture(object_path: str) -> bool:
"""
Delete a user's profile picture
"""
try:
minio_client.remove_object(bucket_name=USER_PFP_BUCKET, object_name=object_path)
return True
except S3Error as e:
print(f"Failed to delete profile picture: {str(e)}")
return False

View File

@@ -0,0 +1,369 @@
import asyncio
import json
from datetime import datetime
from typing import Any, Dict, List
import aio_pika
from aio_pika.abc import AbstractIncomingMessage
from config.logging_config import logging
from config.rabbitmq_config import rabbitmq_manager
from crud.experiment_crud import (
update_simulation_result,
)
from crud.machine_crud import (
update_system_status,
update_systems_offline,
)
from redis.asyncio.client import Redis
from sql_models.models import Instance
from sqlalchemy.ext.asyncio import AsyncSession
## ---------------- HEARTBEAT CONSUMING ------------------------
HEARTBEAT_EXCHANGE = "heartbeat"
HEARTBEAT_QUEUE_PREFIX = "heartbeat_monitor_"
# Redis keys for heartbeat tracking
REDIS_DEVICE_HEARTBEATS = "device:heartbeats" # Sorted set: device_id -> timestamp
REDIS_DEVICE_STATUS = "device:status" # Hash: device_id -> status
REDIS_DEVICE_OFFLINE_HISTORY = (
"device:offline_history" # List: history of offline events
)
# Configuration
DEVICE_TIMEOUT_SECONDS = (
120 # Device considered offline after 120 seconds with no heartbeat
)
OFFLINE_SWEEP_INTERVAL = 10 # Check for offline devices every 30 seconds
async def set_offline(redis: Redis, db: AsyncSession):
logging.info("OFFLINE CHECK STARTED")
while True:
try:
# Calculate cutoff (de vices that haven't sent heartbeat in last 120 seconds)
timeout_seconds = 20
cutoff = datetime.now().timestamp() - timeout_seconds
# Get all offline devices (last heartbeat before cutoff)
offline_device_ids = await redis.zrangebyscore(
"device:heartbeats", min=0, max=cutoff
)
offline_device_ids = [int(device_id) for device_id in offline_device_ids]
# Single database update: set all offline devices to offline
if offline_device_ids:
await update_systems_offline(db, offline_device_ids)
await redis.zremrangebyscore("device:heartbeats", min=0, max=cutoff)
await asyncio.sleep(OFFLINE_SWEEP_INTERVAL)
except Exception as e:
print(f"OFFLINE FAILED: {e}")
await asyncio.sleep(OFFLINE_SWEEP_INTERVAL)
async def process_heartbeat(
device_id: int, status: str, timestamp: datetime, redis: Redis, db: AsyncSession
):
"""
Process heartbeat from a device using Redis sorted sets.
"""
# Convert timestamp to float for Redis score
timestamp_float = timestamp.timestamp()
# Update current device's last heartbeat in sorted set
await redis.zadd("device:heartbeats", {str(device_id): timestamp_float})
# Single database update: set current device to online
await update_system_status(db, device_id, status, timestamp)
async def consume_all_heartbeats(redis: Redis, db: AsyncSession):
"""
Consume heartbeats from ALL clients using a fanout exchange.
Each client publishes to the fanout exchange, and this server consumes all.
"""
try:
channel = await rabbitmq_manager.get_consumer_channel()
if not channel:
raise Exception("Failed to get consumer channel")
# Declare the fanout exchange
exchange = await channel.declare_exchange(
HEARTBEAT_EXCHANGE, type=aio_pika.ExchangeType.FANOUT, durable=True
)
# Create a unique queue for this consumer
# Using a random queue name or fixed name for the main server
queue_name = f"{HEARTBEAT_QUEUE_PREFIX}main_server"
queue = await channel.declare_queue(queue_name, durable=False, auto_delete=True)
# Bind the queue to the fanout exchange
await queue.bind(exchange)
logging.info(
f"Started consuming heartbeats from all clients on exchange '{HEARTBEAT_EXCHANGE}'"
)
asyncio.create_task(set_offline(redis, db))
# Start consuming
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
await handle_heartbeat_message(message, redis, db)
except asyncio.CancelledError:
logging.info("Heartbeat consumer task was cancelled.")
except Exception as e:
logging.info(f"Error in heartbeat consumer: {e}")
raise
async def handle_heartbeat_message(
message: AbstractIncomingMessage, redis: Redis, db: AsyncSession
):
"""Extract device_id and status from heartbeat message body."""
try:
body = message.body.decode()
heartbeat_data = json.loads(body)
# Extract device_id and status from body
device_id = int(heartbeat_data.get("device_id"))
status = heartbeat_data.get("status", "OFFLINE")
timestamp = datetime.fromisoformat(heartbeat_data.get("timestamp"))
if not device_id:
print(f"Received heartbeat without device_id: {heartbeat_data}")
return
await process_heartbeat(device_id, status, timestamp, redis, db)
except json.JSONDecodeError as e:
print(f"Failed to parse heartbeat JSON: {e}, raw body: {message.body}")
except Exception as e:
print(f"Error handling heartbeat message: {e}")
# Example usage in your main server lifespan
async def start_heartbeat_monitoring(redis: Redis, db: AsyncSession):
"""Start consuming heartbeats from all clients."""
return asyncio.create_task(consume_all_heartbeats(redis, db))
##-------- PROGRESS CONSUMING --------------------
PROGRESS_EXCHANGE = "progress_report"
PROGRESS_QUEUE_PREFIX = "progress_consumer_"
async def consume_progress_reports(db: AsyncSession):
"""
Consume progress reports from quantum backend systems.
Each backend publishes to the progress_report exchange, and this server consumes all.
"""
try:
channel = await rabbitmq_manager.get_consumer_channel()
if not channel:
raise Exception("Failed to get consumer channel")
# Declare the exchange (must match the publisher's exchange)
exchange = await channel.declare_exchange(
PROGRESS_EXCHANGE,
type=aio_pika.ExchangeType.DIRECT, # DIRECT matches the publisher
durable=True,
)
# Create a unique queue for this consumer
queue_name = f"{PROGRESS_QUEUE_PREFIX}main_server"
queue = await channel.declare_queue(queue_name, durable=False, auto_delete=True)
# Bind the queue to the exchange
await queue.bind(exchange, routing_key="")
logging.info(
f"Started consuming progress reports from exchange '{PROGRESS_EXCHANGE}'"
)
# Start consuming
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
await handle_progress_message(message, db)
except asyncio.CancelledError:
logging.info("Progress report consumer task was cancelled.")
except Exception as e:
logging.error(f"Error in progress report consumer: {e}")
raise
async def handle_progress_message(message: AbstractIncomingMessage, db: AsyncSession):
"""
Handle incoming progress report messages from quantum backends.
Updates task status in database and optionally caches in Redis.
"""
try:
body = message.body.decode()
if body:
progress_data = json.loads(body)
else:
progress_data = None
# Extract task_id from headers (as sent by publisher)
task_id = message.headers.get("task_id")
if not task_id:
logging.warning(f"Received progress report without task_id")
return
status = str(message.headers.get("status"))
if not status:
logging.warning(f"Received progress report without status")
return
# Extract system_id from the progress data
system_id = message.headers.get("system_id")
if not system_id:
logging.warning(f"Progress report for task {task_id} missing system_id")
return
ended_at = None
if status.upper() == "COMPLETE":
ended_at = datetime.now()
# Update simulation result (only update provided fields)
await update_simulation_result(
db=db,
task_id=int(task_id),
comp_system_id=int(system_id),
status_name=status.upper(),
ended_at=ended_at,
simulation_result_data=progress_data,
)
except json.JSONDecodeError as e:
logging.error(
f"Failed to parse progress report JSON: {e}, raw body: {message.body}"
)
except Exception as e:
logging.error(f"Error handling progress report message: {e}")
import traceback
traceback.print_exc()
# Add this to your main server lifespan or startup function
async def start_progress_monitoring(db: AsyncSession):
"""Start consuming progress reports from quantum backends."""
return asyncio.create_task(consume_progress_reports(db))
##-------- EXPERIMENT PUBLISHING --------------------
TEAM_EXCHANGE_PREFIX = "team_"
async def bind_queues_to_team_exchange(team_id: int, max_qubits: int) -> bool:
"""
Bind a system's queue to team exchange with routing keys for each qubit level.
Called by MAIN SERVER when a system joins a team.
Args:
system_id: The system's unique identifier
team_id: Team ID for the exchange
max_qubits: Maximum qubits this system can handle
"""
try:
channel = await rabbitmq_manager.get_publisher_channel()
if not channel:
raise Exception("Failed to get channel")
# Declare the team exchange (topic exchange for qubit-based routing)
exchange_name = f"{TEAM_EXCHANGE_PREFIX}{team_id}"
exchange = await channel.declare_exchange(
exchange_name, type=aio_pika.ExchangeType.TOPIC, durable=True
)
created_queues = []
for qubits in range(1, max_qubits + 1):
queue_name = f"team_{team_id}.qubits_{qubits}"
# Declare the shared queue
queue = await channel.declare_queue(
queue_name,
durable=True,
arguments={
"x-max-priority": 100, # Allow priorities 0-10
},
)
# Bind queue to exchange with routing key
routing_key = f"qubits.{qubits}"
await queue.bind(exchange, routing_key=routing_key)
created_queues.append(queue_name)
print(
f"Created queue {queue_name} and bound to {exchange_name} with key {routing_key}"
)
print(
f"Setup complete for team {team_id}: {len(created_queues)} queues created"
)
return True
except Exception as e:
print(f"Error binding system to team exchange: {e}")
return False
async def publish_task_to_team(task: Instance):
"""
Publish a task to a team exchange.
The exchange will route to the appropriate shared queue based on qubits needed.
Called by MAIN SERVER when distributing tasks.
Args:
team_id: The team to send the task to
qubits_needed: Number of qubits required (determines routing key)
task_data: The task data to send
priority: Message priority (0-10, higher = more important)
"""
try:
channel = await rabbitmq_manager.get_publisher_channel()
if not channel:
raise Exception("Failed to get publisher channel")
# Declare the team exchange
exchange_name = f"{TEAM_EXCHANGE_PREFIX}{task.experiment.team_id}"
exchange = await channel.declare_exchange(
exchange_name, type=aio_pika.ExchangeType.TOPIC, durable=True
)
# Routing key based on qubits needed
routing_key = f"qubits.{task.qubits_needed}"
# Prepare message
message = {
"task_id": task.id,
"qubits_needed": task.qubits_needed,
"data": task.instance_data,
"timestamp": datetime.now().isoformat(),
}
# Publish to exchange
await exchange.publish(
aio_pika.Message(
body=json.dumps(message).encode(),
delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
),
routing_key=routing_key,
)
print(
f"Task published to exchange {exchange_name} with routing key {routing_key} team"
)
except Exception as e:
print(f"Error publishing task to team: {e}")
raise

View File

@@ -2,6 +2,7 @@ from datetime import datetime
from typing import List, Optional, Tuple
from crud.team_crud import get_user_teams
from pydantic import Json
from sql_models.models import (
Experiment,
ExperimentType,
@@ -51,7 +52,6 @@ async def create_experiment_type(
db: AsyncSession,
name: str,
file_frontend: str,
server_path: str,
file_comp_system: str,
description: Optional[str] = None,
) -> ExperimentType:
@@ -72,7 +72,6 @@ async def create_experiment_type(
name=name,
description=description,
file_frontend=file_frontend,
server_path=server_path,
file_comp_system=file_comp_system,
)
db.add(experiment_type)
@@ -82,6 +81,48 @@ async def create_experiment_type(
return experiment_type
async def update_experiment_type(
db: AsyncSession,
id: int,
name: Optional[str] = None,
file_frontend: Optional[str] = None,
file_comp_system: Optional[str] = None,
description: Optional[str] = None,
) -> ExperimentType:
"""Create a new experiment type"""
# Check if experiment type with this name already exists
existing_result = await db.execute(
select(ExperimentType).where(ExperimentType.id == id)
)
existing = existing_result.scalar_one_or_none()
if not existing:
raise HTTPException(
status_code=400, detail=f"Experiment type with name '{name}' not found"
)
if name:
existing.name = name
if file_frontend:
existing.file_frontend = file_frontend
if file_comp_system:
existing.file_comp_system = file_comp_system
if file_comp_system:
existing.description = description
await db.commit()
await db.refresh(existing)
return existing
async def get_experiment_type_by_id(db: AsyncSession, experiment_type_id: int):
"""Get experiment type by ID"""
result = await db.execute(
select(ExperimentType).where(ExperimentType.id == experiment_type_id)
)
return result.scalar_one_or_none()
# ============= EXPERIMENT CRUD =============
@@ -117,7 +158,9 @@ async def create_experiment(
)
db.add(experiment)
await db.commit()
await db.refresh(experiment, attribute_names=["team", "experiment_type"])
await db.refresh(
experiment, attribute_names=["team", "experiment_type", "instances"]
)
return experiment
@@ -232,7 +275,7 @@ async def delete_experiment(
async def create_instance(
db: AsyncSession,
experiment_id: int,
instance_data_id: int,
instance_data: Json,
name: str,
description: Optional[str] = None,
) -> Instance:
@@ -248,7 +291,7 @@ async def create_instance(
# Create instance
instance = Instance(
experiment_id=experiment_id,
instance_data_id=instance_data_id,
instance_data=instance_data,
name=name,
description=description,
)
@@ -262,9 +305,10 @@ async def create_instance(
async def update_instance(
db: AsyncSession,
instance_id: int,
qubits_needed: Optional[int] = None,
name: Optional[str] = None,
description: Optional[str] = None,
simulation_result_id: Optional[int] = None,
instance_data: Optional[Json] = None,
) -> Instance:
"""Update instance data (cannot update instance_data_id)"""
result = await db.execute(
@@ -285,14 +329,11 @@ async def update_instance(
if description is not None:
instance.description = description
if simulation_result_id is not None:
# Verify simulation result exists
sim_result_result = await db.execute(
select(SimulationResult).where(SimulationResult.id == simulation_result_id)
)
if not sim_result_result.scalar_one_or_none():
raise HTTPException(status_code=404, detail="Simulation result not found")
instance.simulation_result_id = simulation_result_id
if instance_data is not None:
instance.instance_data = instance_data
if qubits_needed is not None:
instance.qubits_needed = qubits_needed
await db.commit()
await db.refresh(instance, attribute_names=["experiment", "simulation_result"])
@@ -385,49 +426,42 @@ async def delete_instance(
async def set_simulation_result(
db: AsyncSession,
comp_system_id: int,
simulation_result_id: int,
comp_system_id: int | None,
instance_id: int,
status_name: str,
started_at: Optional[datetime] = None,
started_at: datetime,
ended_at: Optional[datetime] = None,
simulation_result_data: Json = {},
) -> SimulationResult:
"""Create or update simulation result"""
# Verify computational system exists
from sql_models.models import ComputationalSystem
system_result = await db.execute(
select(ComputationalSystem).where(ComputationalSystem.id == comp_system_id)
)
if not system_result.scalar_one_or_none():
raise HTTPException(status_code=404, detail="Computational system not found")
if comp_system_id:
system_result = await db.execute(
select(ComputationalSystem).where(ComputationalSystem.id == comp_system_id)
)
if not system_result.scalar_one_or_none():
raise HTTPException(
status_code=404, detail="Computational system not found"
)
# Get status
status = await get_simulation_status_by_name(db, status_name)
# Check if simulation result already exists
existing_result = await db.execute(
select(SimulationResult).where(
SimulationResult.simulation_result_id == simulation_result_id
)
)
simulation_result = existing_result.scalar_one_or_none()
instance = await get_single_instance(db, instance_id)
if not instance:
raise HTTPException(404, "Instance not found")
if simulation_result:
# Update existing
simulation_result.comp_system_id = comp_system_id
simulation_result.status_id = status.id
simulation_result.started_at = started_at or simulation_result.started_at
simulation_result.ended_at = ended_at
else:
# Create new
simulation_result = SimulationResult(
comp_system_id=comp_system_id,
simulation_result_id=simulation_result_id,
status_id=status.id,
started_at=started_at or datetime.now(),
ended_at=ended_at,
)
db.add(simulation_result)
simulation_result = SimulationResult(
comp_system_id=comp_system_id,
simulation_result=simulation_result_data,
status_id=status.id,
started_at=started_at or datetime.now(),
ended_at=ended_at,
)
instance.simulation_result = simulation_result
db.add(simulation_result)
await db.commit()
await db.refresh(
@@ -451,3 +485,40 @@ async def get_simulation_result_by_id(
)
)
return result.scalar_one_or_none()
async def update_simulation_result(
db: AsyncSession,
task_id: int,
comp_system_id: Optional[int] = None,
status_name: Optional[str] = None,
ended_at: Optional[datetime] = None,
simulation_result_data: Optional[Json] = None,
):
"""Update simulation result - only updates provided fields. started_at cannot be updated."""
# Get the instance with its simulation result
instance = await get_single_instance(db, task_id)
if not instance:
raise HTTPException(404, "Instance not found")
if not instance.simulation_result:
raise HTTPException(404, "Result not found")
# Update existing simulation result
if comp_system_id is not None:
instance.simulation_result.comp_system_id = comp_system_id
if status_name is not None:
status = await get_simulation_status_by_name(db, status_name)
instance.simulation_result.status_id = status.id
if ended_at is not None:
instance.simulation_result.ended_at = ended_at
if simulation_result_data is not None:
instance.simulation_result.simulation_result = simulation_result_data
await db.commit()
return True

View File

@@ -5,6 +5,7 @@ from sql_models.models import ComputationalSystem, SystemStatus, Team, TeamSyste
from sqlalchemy import and_, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from sqlalchemy.sql import update
from sqlalchemy.sql.expression import func
from starlette.exceptions import HTTPException
@@ -45,7 +46,7 @@ async def create_or_get_computational_system(
return existing
# Get the status (will raise exception if invalid)
status = await get_status_by_name(db, "ONLINE")
status = await get_status_by_name(db, "OFFLINE")
# Create new system
system = ComputationalSystem(
@@ -318,9 +319,7 @@ async def get_system_status(
async def update_system_status(
db: AsyncSession,
system_id: int,
status_name: str,
db: AsyncSession, system_id: int, status_name: str, timestamp: datetime
) -> ComputationalSystem:
# Get the system
@@ -337,9 +336,36 @@ async def update_system_status(
# Update system
system.status_id = status.id
system.last_updated = datetime.now()
system.last_updated = timestamp
await db.commit()
await db.refresh(system, attribute_names=["user", "status"])
return system
async def update_systems_offline(
db: AsyncSession,
system_ids: List[int],
) -> None:
"""
Set multiple systems to offline status in one bulk operation.
Args:
db: Database session
system_ids: List of system IDs to mark as offline
"""
if not system_ids:
return
# Get the offline status object (you already have this function)
offline_status = await get_status_by_name(db, "OFFLINE")
# Bulk update all systems in one query
await db.execute(
update(ComputationalSystem)
.where(ComputationalSystem.id.in_(system_ids))
.values(status_id=offline_status.id)
)
await db.commit()

View File

@@ -53,3 +53,20 @@ async def update_user_profile(
await db.flush()
await db.refresh(user)
return user
async def update_user_profile_picture(
db: AsyncSession, keycloak_id: str, profile_picture_path: Optional[str] = None
) -> User | None:
"""Update user's profile picture path"""
result = await db.execute(select(User).where(User.keycloak_id == keycloak_id))
user = result.scalar_one_or_none()
if not user:
return None
user.profile_picture_path = profile_picture_path
await db.commit()
await db.refresh(user)
return user

View File

@@ -2,24 +2,84 @@ from datetime import datetime
from typing import List, Optional
from pydantic import BaseModel
from pydantic.config import ConfigDict
from pydantic.types import Json
from rest_models.machine_models import (
ComputationalSystemShortData,
)
from rest_models.team_models import (
TeamsShortListResponse,
)
from sqlalchemy.sql.sqltypes import JSON
# ============= INSTANCE MODELS =============
class CreateInstanceRequest(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
experiment_id: int
instance_data: Json
name: str
description: Optional[str] = None
class SimpleInstanceData(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
id: int
instance_data: Json
qubits_needed: int
name: str
description: Optional[str] = None
class UpdateInstanceRequest(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
instance_id: int
qubits_needed: int
name: Optional[str] = None
instance_data: Optional[Json] = None
description: Optional[str] = None
class SimulationResultData(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
id: int
comp_system: ComputationalSystemShortData | None
simulation_result: Json
status: str
started_at: datetime
ended_at: Optional[datetime] = None
class InstanceData(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
instance_id: int
name: str
description: Optional[str] = None
instance_data: Json
qubits_needed: int
simulation_result: SimulationResultData | None
class InstanceListResponse(BaseModel):
instances: List[InstanceData]
cur_page: int
total_instances: int
page_size: int
# ============= EXPERIMENT TYPE MODELS =============
class ExperimentTypeData(BaseModel):
class ExperimentTypeList(BaseModel):
id: int
name: str
description: Optional[str] = None
file_frontend: str
server_path: str
file_comp_system: str
# Pydantic models
class CreateExperimentTypeRequest(BaseModel):
name: str
file_frontend: str
server_path: str
file_comp_system: str
description: Optional[str] = None
@@ -27,9 +87,6 @@ class CreateExperimentTypeResponse(BaseModel):
id: int
name: str
description: Optional[str] = None
file_frontend: str
server_path: str
file_comp_system: str
# ============= EXPERIMENT MODELS =============
@@ -42,40 +99,22 @@ class CreateExperimentRequest(BaseModel):
description: Optional[str] = None
class CreateExperimentResponse(BaseModel):
id: int
team_id: int
experiment_type_id: int
name: str
description: Optional[str] = None
created_at: datetime
class UpdateExperimentRequest(BaseModel):
experiment_id: int
name: Optional[str] = None
description: Optional[str] = None
class UpdateExperimentResponse(BaseModel):
id: int
team_id: int
experiment_type_id: int
name: str
description: Optional[str] = None
created_at: datetime
class ExperimentData(BaseModel):
id: int
team_id: int
team_name: Optional[str] = None
experiment_type_id: int
experiment_type_name: Optional[str] = None
team: TeamsShortListResponse
experiment_type: ExperimentTypeList
name: str
description: Optional[str] = None
created_at: datetime
instances_count: int = 0
instances_count: int
instance_preview: List[SimpleInstanceData]
status: str
class ExperimentListResponse(BaseModel):
@@ -85,87 +124,8 @@ class ExperimentListResponse(BaseModel):
page_size: int
# ============= INSTANCE MODELS =============
class CreateInstanceRequest(BaseModel):
experiment_id: int
instance_data_id: int
name: str
description: Optional[str] = None
class CreateInstanceResponse(BaseModel):
id: int
experiment_id: int
instance_data_id: int
name: str
description: Optional[str] = None
simulation_result_id: Optional[int] = None
class UpdateInstanceRequest(BaseModel):
instance_id: int
name: Optional[str] = None
description: Optional[str] = None
class UpdateInstanceResponse(BaseModel):
id: int
experiment_id: int
instance_data_id: int
name: str
description: Optional[str] = None
simulation_result_id: Optional[int] = None
class InstanceData(BaseModel):
id: int
experiment_id: int
instance_data_id: int
name: str
description: Optional[str] = None
simulation_result_id: Optional[int] = None
simulation_status: Optional[str] = None
simulation_started_at: Optional[datetime] = None
simulation_ended_at: Optional[datetime] = None
class InstanceListResponse(BaseModel):
instances: List[InstanceData]
cur_page: int
total_instances: int
page_size: int
# ============= SIMULATION MODELS =============
class StartExperimentRequest(BaseModel):
instance_id: int
comp_system_id: int
simulation_result_id: int
class StartExperimentResponse(BaseModel):
instance_id: int
simulation_result_id: int
status: str
started_at: datetime
class SetSimulationResultRequest(BaseModel):
comp_system_id: int
simulation_result_id: int
status_name: str
started_at: Optional[datetime] = None
ended_at: Optional[datetime] = None
class SimulationResultData(BaseModel):
id: int
simulation_result_id: int
comp_system_id: int
status: str
started_at: datetime
ended_at: Optional[datetime] = None
experiment_id: int

View File

@@ -72,3 +72,8 @@ class SystemStatusResponse(BaseModel):
class ComputationalSystemDeleteRequest(BaseModel):
system_id: int
class ComputationalSystemShortData(BaseModel):
system_id: int
system_name: str

View File

@@ -6,6 +6,7 @@ from sqlalchemy import Boolean, DateTime, Integer, String
from sqlalchemy.orm import Mapped, mapped_column, relationship
from sqlalchemy.orm.properties import ForeignKey
from sqlalchemy.sql.schema import Column, Table
from sqlalchemy.sql.sqltypes import JSON
from typing_extensions import Optional
@@ -211,7 +212,6 @@ class ExperimentType(Base):
file_frontend: Mapped[str] = mapped_column(
String(255), nullable=False
) # Path to frontend file
server_path: Mapped[str] = mapped_column(String(255), nullable=False) # Server path
file_comp_system: Mapped[str] = mapped_column(
String(255), nullable=False
) # Path to computational system file
@@ -259,11 +259,9 @@ class SimulationResult(Base):
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
comp_system_id: Mapped[int] = mapped_column(
Integer, ForeignKey("computational_systems.id"), nullable=False
Integer, ForeignKey("computational_systems.id"), nullable=True
)
simulation_result_id: Mapped[int] = mapped_column(
Integer, nullable=False
) # FK to another microservice
simulation_result: Mapped[JSON] = mapped_column(JSON, nullable=True)
status_id: Mapped[int] = mapped_column(
Integer, ForeignKey("simulation_statuses.id"), nullable=False
)
@@ -293,10 +291,12 @@ class Instance(Base):
simulation_result_id: Mapped[Optional[int]] = mapped_column(
Integer, ForeignKey("simulation_results.id"), nullable=True
)
instance_data_id: Mapped[int] = mapped_column(
Integer, nullable=False
instance_data: Mapped[JSON] = mapped_column(
JSON, nullable=False
) # FK to another microservice
qubits_needed: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
name: Mapped[str] = mapped_column(String(200), nullable=False)
description: Mapped[Optional[str]] = mapped_column(String(1000), nullable=True)