MeMoModelling/gene-expression-bootstrapper
0
1import asyncio2import io3import json4import os5import sys6import tempfile7import uuid8import zipfile9from contextlib import redirect_stdout10from pathlib import Path11 12from fastapi import FastAPI, File, Form, UploadFile13from fastapi.responses import FileResponse, JSONResponse, StreamingResponse14from fastapi.staticfiles import StaticFiles15 16app = FastAPI()17 18# In-memory job store19jobs: dict = {}20 21 22@app.post("/api/process")23async def process(24 model_files: list[UploadFile] = File(...),25 mapping_files: list[UploadFile] = File(...),26 expr_file: UploadFile = File(...),27 batch_count: int = Form(1000),28):29 job_id = str(uuid.uuid4())30 jobs[job_id] = {31 "status": "running",32 "progress": 0,33 "messages": [],34 "output_zip": None,35 "error": None,36 "file_count": batch_count,37 "species_count": len(model_files),38 }39 40 # Save uploaded files to temp dir41 tmp = tempfile.mkdtemp()42 43 model_paths = []44 for f in model_files:45 content = await f.read()46 path = os.path.join(tmp, f.filename)47 with open(path, "wb") as out:48 out.write(content)49 model_paths.append(path)50 51 mapping_paths = []52 for f in mapping_files:53 content = await f.read()54 path = os.path.join(tmp, f.filename)55 with open(path, "wb") as out:56 out.write(content)57 mapping_paths.append(path)58 59 expr_content = await expr_file.read()60 expr_path = os.path.join(tmp, expr_file.filename)61 with open(expr_path, "wb") as out:62 out.write(expr_content)63 64 species_prefixes = [Path(f.filename).stem.split("_")[0] for f in model_files]65 66 asyncio.create_task(67 run_bootstrap(job_id, model_paths, mapping_paths, expr_path, species_prefixes, batch_count, tmp)68 )69 70 return {"job_id": job_id}71 72 73async def run_bootstrap(job_id, model_paths, mapping_paths, expr_path, species_prefixes, batch_count, tmp):74 job = jobs[job_id]75 76 def add_msg(text, type_=""):77 job["messages"].append({"text": text, "type": type_})78 79 try:80 add_msg(f"Starting bootstrap for: {', '.join(species_prefixes)}", "info")81 job["progress"] = 1082 83 output_dir = os.path.join(tmp, "output")84 os.makedirs(output_dir, exist_ok=True)85 86 captured = io.StringIO()87 88 def run_sync():89 sys.path.insert(0, "/app")90 from utils.bootstrap_genes import bootstrap_genes91 with redirect_stdout(captured):92 bootstrap_genes(93 model_pre_filenames=model_paths,94 mapping_filenames=mapping_paths,95 species_prefixes=species_prefixes,96 combined_geneExpr_filename=expr_path,97 geneExpr_folder=output_dir,98 batch_count=batch_count,99 )100 101 job["progress"] = 20102 loop = asyncio.get_event_loop()103 await loop.run_in_executor(None, run_sync)104 job["progress"] = 80105 106 # Relay only key status lines — skip all Warning lines107 for line in captured.getvalue().strip().split("\n"):108 line = line.strip()109 if not line or line.startswith("Warning"):110 continue111 if "Write to" in line:112 add_msg("Writing output files…", "ok")113 elif line.startswith("Read") or line.startswith("Species") or line.startswith("Bootstrap"):114 add_msg(line, "info")115 116 add_msg("Compressing output files…", "info")117 job["progress"] = 90118 119 zip_path = os.path.join(tmp, "geneExpr_bootstrapped.zip")120 output_files = sorted([f for f in os.listdir(output_dir) if f.endswith(".csv")])121 with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:122 for fname in output_files:123 zf.write(os.path.join(output_dir, fname), arcname=fname)124 125 job["output_zip"] = zip_path126 job["file_count"] = batch_count127 job["species_count"] = len(species_prefixes)128 job["progress"] = 100129 job["status"] = "done"130 add_msg(f"Done! {batch_count} files generated.", "ok")131 132 except Exception as e:133 import traceback134 job["status"] = "error"135 job["error"] = str(e)136 add_msg(f"Error: {e}", "err")137 print(traceback.format_exc(), flush=True)138 139 140@app.get("/api/progress/{job_id}")141async def progress_stream(job_id: str):142 async def event_stream():143 last_idx = 0144 while True:145 job = jobs.get(job_id)146 if not job:147 yield f"data: {json.dumps({'status': 'error', 'error': 'Job not found'})}\n\n"148 break149 150 new_msgs = job["messages"][last_idx:]151 last_idx = len(job["messages"])152 153 for msg in new_msgs:154 payload = {155 "message": msg["text"],156 "type": msg["type"],157 "progress": job["progress"],158 "label": msg["text"],159 }160 yield f"data: {json.dumps(payload)}\n\n"161 162 if job["status"] == "done":163 yield f"data: {json.dumps({'status': 'done', 'progress': 100, 'file_count': job.get('file_count'), 'species_count': job.get('species_count')})}\n\n"164 break165 elif job["status"] == "error":166 yield f"data: {json.dumps({'status': 'error', 'error': job.get('error', 'Unknown error')})}\n\n"167 break168 169 if not new_msgs:170 yield f"data: {json.dumps({'progress': job['progress'], 'label': 'Processing…'})}\n\n"171 172 await asyncio.sleep(1)173 174 return StreamingResponse(175 event_stream(),176 media_type="text/event-stream",177 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},178 )179 180 181@app.get("/api/download/{job_id}")182async def download(job_id: str):183 job = jobs.get(job_id, {})184 zip_path = job.get("output_zip")185 if not zip_path or not os.path.exists(zip_path):186 return JSONResponse({"error": "File not found"}, status_code=404)187 return FileResponse(188 zip_path,189 filename="geneExpr_bootstrapped.zip",190 media_type="application/zip",191 )192 193 194# Serve frontend — must be last195app.mount("/", StaticFiles(directory="/app/frontend", html=True), name="frontend")196