app.py

import os
import threading
import time
import uuid

import numpy as np
from flask import Flask, jsonify, render_template, request

from pipeline import run_pipeline, _get_esm, _get_phlam, _get_annotation_retriever, _get_lifestyle_host
from utils.genome_neighbor_retrieval import GenomeNeighborRetriever
from utils.gene_context_retrieval import get_gene_context_retriever as _get_gene_context_retriever
from utils.synteny import build_synteny, neighbor_track, query_track
from utils.lazy import lazy_singleton

app = Flask(__name__)

# Background neighbor-search jobs: the imgvr_inphared indexes are large
# enough (GB-scale, network-mounted) that a cold search can take tens of
# seconds, so /analyze starts them and returns immediately; the frontend
# polls /job/<id> (same pattern as phlam-demo's HOMOLOGY_JOBS/NEIGHBOR_JOBS).
JOBS = {}
_JOBS_LOCK = threading.Lock()

# Query genome data kept for a later /synteny call (embeddings aren't sent
# back to the client, so they're cached here keyed by analysis_id).
ANALYSES = {}

# _get_gene_context_retriever is the SAME shared singleton utils/synteny.py's
# neighbor_track uses -- one index load, not a second independent copy.
# _get_genome_retriever has only one consumer (this file), wrapped here.
_get_genome_retriever = lazy_singleton(GenomeNeighborRetriever)


def _prune(store, keep=8):
    if len(store) <= keep:
        return
    stale = sorted(store, key=lambda k: store[k]["created"])[:-keep]
    for k in stale:
        store.pop(k, None)


def _run_job(job_id, fn):
    job = JOBS.get(job_id)
    if job is None:
        return
    try:
        job["result"] = fn()
        job["status"] = "done"
    except Exception as e:
        print(f"[job {job_id}] failed: {e}")
        job["status"] = "error"
        job["error"] = str(e)


def _start_job(fn):
    job_id = uuid.uuid4().hex
    JOBS[job_id] = {"status": "pending", "result": None, "error": None, "created": time.time()}
    with _JOBS_LOCK:
        _prune(JOBS)
    threading.Thread(target=_run_job, args=(job_id, fn), daemon=True).start()
    return job_id


def _warm_up(name, fn):
    """Touches one model/index at process start so the first real request
    isn't the one paying the cold-load cost -- especially the big FAISS
    indexes on a network-mounted bucket. Each runs in its own thread so
    one slow index doesn't delay the others."""
    try:
        t0 = time.time()
        fn()
        print(f"[warmup] {name} ready in {time.time() - t0:.1f}s")
    except Exception as e:
        print(f"[warmup] {name} failed: {e}")


for _name, _fn in (
    ("esm", _get_esm), ("phlam", _get_phlam), ("annotation", _get_annotation_retriever),
    ("lifestyle_host", _get_lifestyle_host),
    ("genome_retriever", _get_genome_retriever), ("gene_context_retriever", _get_gene_context_retriever),
):
    threading.Thread(target=_warm_up, args=(_name, _fn), daemon=True).start()


@app.route("/", methods=["GET"])
def index():
    return render_template("index.html")


@app.route("/analyze", methods=["POST"])
def analyze():
    data = request.get_json(silent=True) or {}
    raw_input = data.get("input", "")
    try:
        result = run_pipeline(raw_input)
    except ValueError as e:
        return jsonify({"error": str(e)}), 400
    except Exception as e:
        print(f"[analyze] unexpected failure: {e}")
        return jsonify({"error": "Analysis failed -- check server logs."}), 500

    esm_embeds = result.pop("esm_embeds")
    phlam_protein = result.pop("phlam_protein")
    phlam_genome = np.asarray(result["phlam_genome"], dtype=np.float32)[None, :]

    analysis_id = uuid.uuid4().hex
    ANALYSES[analysis_id] = {
        "genome_name": result["genome_name"], "genome_length": result["genome_length"], "genes": result["genes"],
        "esm_embeds": esm_embeds, "phlam_protein": phlam_protein, "created": time.time(),
    }
    with _JOBS_LOCK:
        _prune(ANALYSES)

    jobs = {
        "genome_neighbors": _start_job(lambda: _get_genome_retriever().query(phlam_genome)[0]),
        "gene_context_neighbors": _start_job(lambda: _get_gene_context_retriever().query(phlam_protein)),
        "annotations": _start_job(lambda: _get_annotation_retriever().query(esm_embeds)),
    }
    result["analysis_id"] = analysis_id
    result["jobs"] = jobs
    return jsonify(result)


@app.route("/synteny", methods=["POST"])
def synteny():
    data = request.get_json(silent=True) or {}
    analysis = ANALYSES.get(data.get("analysis_id"))
    if analysis is None:
        return jsonify({"error": "analysis expired -- re-run the analysis"}), 404
    genome_ids = [int(g) for g in data.get("genome_ids", [])][:11]  # query + up to 11 neighbors

    def run():
        tracks = [query_track(analysis["genome_name"], analysis["genome_length"], analysis["genes"],
                               analysis["esm_embeds"], analysis["phlam_protein"])]
        tracks += [neighbor_track(gid) for gid in genome_ids]
        return build_synteny(tracks)

    return jsonify({"job": _start_job(run)})


@app.route("/job/<job_id>", methods=["GET"])
def job_status(job_id):
    job = JOBS.get(job_id)
    if job is None:
        return jsonify({"status": "unknown"}), 404
    return jsonify({"status": job["status"], "result": job["result"], "error": job["error"]})


if __name__ == "__main__":
    port = int(os.environ.get("PORT", "7860"))
    debug = os.environ.get("FLASK_DEBUG", "0").lower() in ("1", "true", "yes")
    app.run(debug=debug, host="0.0.0.0", port=port, threaded=True)