9e875c5a04
Added `type` to experiment runs in `sample.py` and improved filtering in `processing.py` to match experiments by both `sample_id` and `run_id`. Removed extensive unnecessary code in `testfunctions.ipynb` for clarity and maintenance.
96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
import json
|
|
import asyncio
|
|
from fastapi import APIRouter, Depends
|
|
from fastapi.encoders import jsonable_encoder
|
|
from sqlalchemy.orm import Session
|
|
from starlette.responses import StreamingResponse
|
|
from app.models import (
|
|
JobStatus,
|
|
Jobs as JobModel,
|
|
ExperimentParameters as ExperimentParametersModel,
|
|
Sample as SampleModel,
|
|
)
|
|
from app.schemas import JobsResponse, JobsUpdate
|
|
from app.dependencies import get_db
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
async def job_event_generator(db: Session):
|
|
while True:
|
|
jobs = db.query(JobModel).filter(JobModel.status == JobStatus.TODO).all()
|
|
job_items = []
|
|
for job in jobs:
|
|
sample = db.query(SampleModel).filter_by(id=job.sample_id).first()
|
|
experiment = (
|
|
db.query(ExperimentParametersModel)
|
|
.filter(
|
|
ExperimentParametersModel.sample_id == sample.id,
|
|
ExperimentParametersModel.id == job.run_id,
|
|
)
|
|
.first()
|
|
)
|
|
|
|
job_item = JobsResponse(
|
|
job_id=job.id,
|
|
sample_id=sample.id,
|
|
run_id=job.run_id,
|
|
sample_name=sample.sample_name,
|
|
status=job.status,
|
|
type=experiment.type,
|
|
created_at=job.created_at,
|
|
updated_at=job.updated_at,
|
|
data_collection_parameters=sample.data_collection_parameters,
|
|
experiment_parameters=experiment.beamline_parameters
|
|
if experiment
|
|
else None,
|
|
filepath=experiment.dataset.get("filepath")
|
|
if experiment and experiment.dataset
|
|
else None,
|
|
slurm_id=job.slurm_id,
|
|
)
|
|
|
|
job_items.append(job_item)
|
|
|
|
if job_items:
|
|
# Use Pydantic's .json() for each item, if you need a fine structure, or:
|
|
serialized = jsonable_encoder(job_items)
|
|
yield f"data: {json.dumps(serialized)}\n\n"
|
|
|
|
await asyncio.sleep(5)
|
|
|
|
|
|
# A reasonable heartbeat/refresh
|
|
|
|
|
|
@router.get("/jobs/stream")
|
|
async def stream_jobs(db: Session = Depends(get_db)):
|
|
return StreamingResponse(job_event_generator(db), media_type="text/event-stream")
|
|
|
|
|
|
@router.post(
|
|
"/jobs/update_status", response_model=JobsUpdate, operation_id="update_status"
|
|
)
|
|
def update_jobs_status(payload: JobsUpdate, db: Session = Depends(get_db)):
|
|
# Fetch the job by job_id
|
|
job = db.query(JobModel).filter(JobModel.id == payload.job_id).first()
|
|
if not job:
|
|
# Optionally, use HTTPException for proper status code
|
|
from fastapi import HTTPException
|
|
|
|
raise HTTPException(status_code=404, detail="Job not found")
|
|
|
|
# Update the status
|
|
job.status = payload.status
|
|
job.slurm_id = payload.slurm_id
|
|
# Optionally update 'updated_at'
|
|
from datetime import datetime
|
|
|
|
job.updated_at = datetime.now()
|
|
|
|
db.commit()
|
|
db.refresh(job)
|
|
|
|
# Return the updated job's info as response
|
|
return JobsUpdate(job_id=job.id, status=job.status, slurm_id=job.slurm_id)
|