"""Offline CPU inference for Leak CURBER enzyme classification, version 2. Pipeline -------- 1. Safety net: a quick DIAMOND search, so every input has homology evidence early on. 2. Main search: DIAMOND `--very-sensitive` with full alignments, read through a pipe query by query and turned into features on arrival, including EC-specific residue profiles. 3. ESM-2 and the hierarchical head, in a child process, for queries with weak homology. 4. A bag of LightGBM LambdaRank models ranks each (query, candidate EC) pair. 5. Anything with no evidence falls back to the most frequent training EC. Running inside the rules ------------------------ The task says no file written by predict.py may exceed 16 MiB, and a limit like that may be enforced by the OS, killing the process at the first large write. So the DIAMOND index is built when the image is built (see Dockerfile), `--hit-membuf` keeps DIAMOND's intermediate hits in memory, and results arrive through stdout. Nothing large touches the disk. Every stage has a deadline and a failure path. A slow or failing stage costs accuracy on the queries it did not reach; it never costs the run, because a valid output file is always written. """ import argparse import gzip import json import os import shutil import subprocess import sys import tempfile import threading import time from pathlib import Path import lightgbm as lgb # imported before anything that could pull in torch; see embed_worker.py import numpy as np import pandas as pd sys.path.insert(0, str(Path(__file__).parent)) from features import OUTFMT, assemble, query_features, read_hits # noqa: E402 from inputs import prepare_inputs # noqa: E402 ROOT = Path(__file__).parent WEIGHTS = ROOT / "weights" T0 = time.time() # The evaluator allows 600 s. Evidence gathering stops at BUDGET; the rest is for ranking and writing. BUDGET = float(os.environ.get("EC_TIME_BUDGET", 540)) THREADS = int(os.environ.get("EC_THREADS", 2)) def elapsed(): return time.time() - T0 def log(msg): print(f"[{elapsed():6.1f}s] {msg}", flush=True) def diamond_bin(): return shutil.which("diamond") or str(ROOT / "bin" / "diamond") def database(tmp): """The image carries a prebuilt index. Building here is only for runs outside the image.""" built = Path(os.environ.get("EC_DMND", WEIGHTS / "train.dmnd")) if built.exists(): return built.with_suffix("") fa = Path(tmp) / "db.fa" with gzip.open(WEIGHTS / "train_db.fa.gz", "rb") as src, open(fa, "wb") as dst: shutil.copyfileobj(src, dst) db = Path(tmp) / "db" subprocess.run([diamond_bin(), "makedb", "--in", str(fa), "-d", str(db), "--threads", str(THREADS), "--quiet"], check=True) fa.unlink() return db def blast(query, db, tmp, columns, k, mode, evalue): cmd = [diamond_bin(), "blastp", "-q", str(query), "-d", str(db), "--outfmt", "6", *columns, "-k", str(k), "--max-hsps", "1", "--evalue", str(evalue), "--threads", str(THREADS), "--quiet", "--tmpdir", str(tmp), "--hit-membuf"] if mode: cmd.append(f"--{mode}") return subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1 << 20) def stream(proc, deadline, k, keep_aln, on_query): """Feed hits to `on_query` until DIAMOND ends or the deadline kills it. Partial results stay.""" timer = threading.Timer(max(deadline - elapsed(), 0.1), proc.kill) timer.start() done = 0 try: for q, rows in read_hits(proc.stdout, k, keep_aln): on_query(q, rows) done += 1 finally: timer.cancel() proc.stdout.close() proc.wait() return done, proc.returncode def embed(seqs, order, cands, cfg, tmp, budget): """Head evidence from an isolated child process; see embed_worker.py for why.""" job = Path(tmp) / "emb_job.json" res = Path(tmp) / "emb_out.json" job.write_text(json.dumps({"seqs": {str(i): seqs[i] for i in order}, "order": [int(i) for i in order], "cands": {str(i): cands[i] for i in order}, "cfg": cfg, "weights": str(WEIGHTS), "budget": budget, "threads": THREADS})) try: proc = subprocess.run([sys.executable, str(ROOT / "embed_worker.py"), str(job), str(res)], capture_output=True, text=True, timeout=budget + 15) except subprocess.TimeoutExpired: log("embedding worker exceeded its time limit") return {} if proc.returncode != 0 or not res.exists(): log(f"embedding worker failed (rc={proc.returncode}): {proc.stderr.strip()[-300:]}") return {} log(proc.stderr.strip() or "embedding worker finished") raw = json.loads(res.read_text()) return {int(k): {ec: tuple(v) for ec, v in d.items()} for k, d in raw.items()} def rank(F, cfg): scores = np.zeros(len(F)) for name in cfg["rerankers"]: booster = lgb.Booster(model_file=str(WEIGHTS / name)) p = booster.predict(F, num_threads=THREADS) scores += (p - p.mean()) / (p.std() + 1e-9) return scores def evidence_and_rank(seqs, cfg, db_ecs, ec_count, ec_prior, vocab, tmp): n = len(seqs) scanned = [None] * n fa = Path(tmp) / "q.fa" with open(fa, "w") as f: for i, s in enumerate(seqs): f.write(f">{i}\n{s}\n") db = database(tmp) log("DIAMOND database ready") def on_query(q, rows): scanned[q] = query_features(rows, seqs[q], db_ecs, ec_count, ec_prior, cfg) plain = OUTFMT[:9] try: proc = blast(fa, db, tmp, plain, cfg["k_quick"], "fast", cfg["evalue"]) got, rc = stream(proc, BUDGET * 0.35, cfg["k_quick"], 0, on_query) log(f"safety-net search covered {got}/{n} queries (rc={rc})") except Exception as exc: log(f"safety-net search failed: {exc}") try: proc = blast(fa, db, tmp, OUTFMT, cfg["k"], cfg["sensitivity"], cfg["evalue"]) got, rc = stream(proc, BUDGET - cfg["reserve"], cfg["k"], cfg["prof_hits"], on_query) log(f"main search covered {got}/{n} queries (rc={rc})") except Exception as exc: log(f"main search failed: {exc}") # Least evidence first, then shortest: the queries that gain most for the least time. best = np.array([r[2]["q_level"][0] if r else 0.0 for r in scanned]) weak = [i for i in np.where(best < cfg["emb_bits_gate"])[0].tolist()] weak.sort(key=lambda i: (best[i] >= 80, len(seqs[i]))) log(f"{len(weak)} queries below the embedding gate") emb = {} room = BUDGET - 20 - elapsed() if weak and room > 30: try: emb = embed(seqs, weak, {i: (scanned[i][0] if scanned[i] else []) for i in weak}, cfg, tmp, room) except Exception as exc: # embedding is extra evidence; homology still stands log(f"embedding stage skipped: {exc}") q, ec, F = assemble(scanned, n, emb, ec_prior, ec_count, seqs, cfg["extra_from_emb"]) log(f"{len(q)} candidate rows") scores = rank(F, cfg) if len(F) else np.zeros(0) best_ec = {} for qi, e, sc in zip(q.tolist(), ec.tolist(), scores.tolist()): if e in vocab and (qi not in best_ec or sc > best_ec[qi][1]): best_ec[qi] = (e, sc) log(f"ranked predictions for {len(best_ec)} queries") return best_ec def predict_file(input_path, output_path): cfg = json.loads((WEIGHTS / "config.json").read_text()) vocab = set(json.loads((ROOT / "ec_vocabulary.json").read_text())) frame = prepare_inputs(pd.read_parquet(input_path)) seqs = frame.sequence.tolist() log(f"{len(seqs)} unique inputs") db_ecs = json.loads((WEIGHTS / "db_ecs.json").read_text()) ec_prior = json.loads((WEIGHTS / "ec_prior.json").read_text()) ec_count = json.loads((WEIGHTS / "ec_count.json").read_text()) fallback = [e for e in json.loads((WEIGHTS / "prior.json").read_text()) if e in vocab][:1] tmp = tempfile.mkdtemp(dir="/tmp") best_ec = {} try: best_ec = evidence_and_rank(seqs, cfg, db_ecs, ec_count, ec_prior, vocab, tmp) except Exception as exc: # a valid, weaker output beats a failed evaluation log(f"pipeline failed, writing fallback predictions: {exc!r}") finally: shutil.rmtree(tmp, ignore_errors=True) with open(output_path, "w") as f: for i, key in enumerate(frame.id): pred = [best_ec[i][0]] if i in best_ec else fallback f.write(json.dumps({"id": key, "prediction": pred}) + "\n") log(f"wrote {len(seqs)} predictions to {output_path}") if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", required=True) parser.add_argument("--output", required=True) args = parser.parse_args() predict_file(args.input, args.output)