"""Offline CPU inference: DIAMOND homology search against train + UniProt references, bitscore vote + learned re-ranker (mode.json). 1. Search every unique input against the training set (diamond blastp --very-sensitive, 50 hits). 2. Re-search inputs with no hit or a weak best hit (<=30% identity) with --ultra-sensitive when time allows. 3. Collect candidate ECs from the hits, compute features and score them with a gradient-boosted tree ensemble (pure numpy). Best EC first. Safety: environment diagnostics are logged; DIAMOND is smoke-tested and retried with safer settings; if it cannot run at all, a pure numpy/scipy 3-mer nearest-neighbour fallback produces the predictions. """ import argparse import gzip import json import os import platform import resource import shutil import signal import subprocess import sys import tarfile import tempfile import threading import time from pathlib import Path import numpy as np import pandas as pd from inputs import prepare_inputs from features import query_candidates, parse_line from ranker import TreeRanker ROOT = Path(__file__).resolve().parent TOP_K = 50 WEAK_PID = 30.0 def load_mode(): cfg = {'scoring': 'learned', 'vote_k': 10, 'vote_power': 4, 'weak_pid': WEAK_PID, 'weak_max': 3000, 'ultra_guard': 215} try: cfg.update(json.loads((Path(__file__).resolve().parent / 'weights' / 'mode.json').read_text())) except Exception: pass return cfg MAX_CHUNK_BYTES = 12 * 1024 * 1024 # every file we write stays under 16 MiB MAX_CHUNK_SEQS = 8000 # bounds DIAMOND's temporary seed-hit files too OUTFMT = ['qseqid', 'sseqid', 'pident', 'length', 'qlen', 'slen', 'evalue', 'bitscore', 'qstart', 'qend', 'sstart', 'send'] START = time.time() # DIAMOND settings to try, most efficient first. CONFIGS = [(['-c1', '-b4'], 2), (['-c4', '-b1'], 1)] # (extra args, threads) def log(msg): print(f'[{time.time() - START:7.1f}s] {msg}', file=sys.stderr, flush=True) def diagnostics(): def probe(fn): try: fn() except Exception as exc: # diagnostics must never break inference log(f'diagnostics: {exc!r}') probe(lambda: log(f'python {sys.version.split()[0]} | {platform.platform()} | machine={platform.machine()} | uid={os.getuid()}')) probe(lambda: log(f'cpu_count={os.cpu_count()} affinity={len(os.sched_getaffinity(0))}')) probe(lambda: log(' '.join(f'{n}={resource.getrlimit(getattr(resource, n))}' for n in ('RLIMIT_FSIZE', 'RLIMIT_AS', 'RLIMIT_DATA', 'RLIMIT_STACK', 'RLIMIT_NPROC', 'RLIMIT_NOFILE')))) def cpu(): info = Path('/proc/cpuinfo').read_text() model = next((l.split(':', 1)[1].strip() for l in info.splitlines() if l.startswith('model name')), '?') flags = next((l.split(':', 1)[1].split() for l in info.splitlines() if l.startswith('flags')), []) log(f'cpu={model} | ' + ' '.join(f'{f}={f in flags}' for f in ('sse4_1', 'avx2', 'avx512f'))) probe(cpu) probe(lambda: log(f'/tmp free={shutil.disk_usage("/tmp").free / 1e9:.2f} GB')) for p in ('/sys/fs/cgroup/memory.max', '/sys/fs/cgroup/pids.max'): probe(lambda p=p: log(f'{p}={Path(p).read_text().strip()}') if Path(p).exists() else None) def cpu_probe(): """Log real usable compute: matmul GFLOP/s with 1, 2, 4, 8, 16 threads (~5 s). Changes nothing.""" try: import numpy as _np from concurrent.futures import ThreadPoolExecutor a = _np.random.rand(768, 768).astype(_np.float32) def work(_): for _i in range(8): a @ a out = [] for t in (1, 2, 4, 8, 16): t0 = time.time() with ThreadPoolExecutor(t) as ex: list(ex.map(work, range(t))) dt = time.time() - t0 out.append(f'{t}t={t * 8 * 2 * 768 ** 3 / dt / 1e9:.0f}') log('cpu probe GFLOP/s: ' + ' '.join(out)) except Exception as exc: log(f'cpu probe error: {exc!r}') def diamond_binaries(): """Official release binary (runtime SIMD dispatch) first, then our generic static build.""" found = [] official = ROOT / 'bin' / 'diamond' if not (official.exists() and os.access(official, os.X_OK)): try: out = Path('/tmp/diamond_bin') out.mkdir(parents=True, exist_ok=True) official = out / 'diamond' with tarfile.open(ROOT / 'bin' / 'diamond-linux64.tar.gz') as tar: official.write_bytes(tar.extractfile('diamond').read()) official.chmod(0o755) except Exception as exc: log(f'cannot extract official diamond: {exc!r}') official = None if official: found.append(('official', str(official))) generic = ROOT / 'bin' / 'diamond-generic' if generic.exists(): if not os.access(generic, os.X_OK): tmp = Path('/tmp/diamond_bin/diamond-generic') tmp.parent.mkdir(parents=True, exist_ok=True) shutil.copyfile(generic, tmp) tmp.chmod(0o755) generic = tmp found.append(('generic', str(generic))) return found def write_chunks(items, workdir, tag): """items: iterable of (query index, sequence).""" paths, buf, size = [], [], 0 def flush(): nonlocal buf, size if buf: p = Path(workdir) / f'{tag}{len(paths)}.fa' p.write_text(''.join(buf)) paths.append(p) buf, size = [], 0 for i, seq in items: rec = f'>{i}\n{seq}\n' if size + len(rec) > MAX_CHUNK_BYTES or len(buf) >= MAX_CHUNK_SEQS: flush() buf.append(rec) size += len(rec) flush() return paths def describe_rc(rc): if rc < 0: try: return f'killed by signal {signal.Signals(-rc).name}' except ValueError: return f'killed by signal {-rc}' return f'exit code {rc}' DEADLINE = 420.0 # seconds after start; DIAMOND is killed past this point def diamond_bins(): """'--bin N' spreads temporary seed hits over N files: more bins keep each file under the 16 MiB size cap, but every bin is an open file and the sandbox allows only ~256 (RLIMIT_NOFILE).""" try: soft = resource.getrlimit(resource.RLIMIT_NOFILE)[0] if soft == resource.RLIM_INFINITY: soft = 4096 except Exception: soft = 256 return int(max(16, min(128, (soft - 64) // 1))) def run_diamond(chunks, mode, cfg, workdir, depth=0, bins=None, db='train.dmnd'): """Search each FASTA chunk. The evaluator caps every written file at 16 MiB (RLIMIT_FSIZE); '--bin 256' spreads DIAMOND's temporary seed hits over many small files, and a chunk that still hits the cap (SIGXFSZ) is split into smaller chunks and retried.""" (name, binary), extra, threads = cfg bins = bins or diamond_bins() hits = {} for chunk in chunks: errf = Path(workdir) / 'diamond.err' cmd = [binary, 'blastp', '-q', str(chunk), '-d', str(ROOT / 'weights' / db), '-k', str(TOP_K), '--threads', str(threads), '--quiet', *extra, '--bin', str(bins), '--tmpdir', str(workdir), f'--{mode}', '--outfmt', '6', *OUTFMT] local = {} with errf.open('w') as err: proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=err, text=True) watchdog = threading.Timer(max(1.0, DEADLINE - (time.time() - START)), proc.kill) watchdog.start() for line in proc.stdout: q, h = parse_line(line) local.setdefault(int(q), []).append(h) rc = proc.wait() watchdog.cancel() if time.time() - START >= DEADLINE - 1 and rc != 0: raise RuntimeError(f'diamond[{name}] {mode} stopped by watchdog at {DEADLINE:.0f}s') if rc == -signal.SIGXFSZ and depth < 3: recs = chunk.read_text().split('>')[1:] if len(recs) > 1: log(f'diamond[{name}] hit the file-size cap on {len(recs)} queries; splitting into 4') parts = [] for k in range(4): sub = recs[k::4] if sub: p = chunk.with_name(f'{chunk.stem}_{depth}_{k}.fa') p.write_text(''.join('>' + r for r in sub)) parts.append(p) chunk.unlink() hits.update(run_diamond(parts, mode, cfg, workdir, depth + 1, bins, db)) continue msg = errf.read_text()[-2000:].strip() if (rc != 0 and errf.exists()) else '' if rc != 0 and 'Too many open files' in msg and bins > 16: log(f'diamond[{name}] too many open files with --bin {bins}; retrying with {bins // 2}') hits.update(run_diamond([chunk], mode, cfg, workdir, depth, bins // 2, db)) continue if rc != 0: raise RuntimeError(f'diamond[{name}] {mode} {extra} t={threads}: {describe_rc(rc)}; stderr: {msg or "(empty)"}') hits.update(local) chunk.unlink() return hits class EcReliability: def __init__(self, path): with gzip.open(path, 'rt') as f: t = json.load(f) self.t, self.A = t, t['A'] def _enc(self, table, key, prior): s, c = self.t[table].get(key, (0.0, 0)) return (s + self.A * prior) / (c + self.A), np.log1p(c) def matrix(self, cands): rows = [] for e in cands: a, an = self._enc('ec', e, self.t['prior']) b, bn = self._enc('ec_top', e, self.t['prior_top']) c, cn = self._enc('l3', e.rsplit('.', 1)[0], self.t['prior']) rows.append([a, an, b, bn, c, cn]) return np.array(rows, dtype=np.float64) def homology_predictions(seqs, labels, freq): """Returns {query index: ranked EC list}, or None if DIAMOND cannot run.""" n = len(seqs) main_mode = 'very-sensitive' if n <= 14000 else ('sensitive' if n <= 40000 else 'fast') # v14: larger reference workdir = tempfile.mkdtemp(dir='/tmp') try: binaries = diamond_binaries() for name, b in binaries: try: log(f'{name}: ' + subprocess.run([b, 'version'], capture_output=True, text=True, timeout=60).stdout.strip()) except Exception as exc: log(f'{name}: version check failed: {exc!r}') sample = [(i, seqs[i]) for i in range(min(20, len(seqs)))] hits, cfg = None, None for b in binaries: for extra, threads in CONFIGS: c = (b, extra, threads) try: run_diamond(write_chunks(sample, workdir, 's'), 'fast', c, workdir) log(f'smoke test ok: {b[0]} {extra} t={threads}') except Exception as exc: log(f'smoke test failed: {exc}') continue t0 = time.time() mode = main_mode if b[0] == 'official' or main_mode == 'fast' else 'sensitive' try: hits = run_diamond(write_chunks(enumerate(seqs), workdir, 'a'), mode, c, workdir) cfg = c break except Exception as exc: log(f'main search failed after {time.time() - t0:.0f}s: {exc}') if time.time() - START > 240: break if hits is not None or time.time() - START > 240: break if hits is None: return None t1 = time.time() - START mcfg = load_mode() weak = [i for i in range(n) if i not in hits or hits[i][0][1] <= float(mcfg['weak_pid'])] log(f'main search ({mode}, {cfg[0][0]}) done: {len(hits)}/{n} with hits, {len(weak)} weak') # The second pass costs about as much as the first, so run it only with a wide margin. if weak and len(weak) <= int(mcfg['weak_max']) and t1 < float(mcfg['ultra_guard']) and cfg[0][0] == 'official': try: extra = run_diamond(write_chunks(((i, seqs[i]) for i in weak), workdir, 'b'), 'ultra-sensitive', cfg, workdir) hits.update(extra) log(f'ultra-sensitive pass done: {len(extra)} re-scored') except Exception as exc: log(f'ultra-sensitive pass skipped: {exc}') # v16: optional pass over a second, larger reference (unreviewed TrEMBL, transcript-level evidence) # in DIAMOND's fast mode. It only adds hits; if time is short or it fails, it is skipped. xdb = mcfg.get('extra_db') t2 = time.time() - START if xdb and (ROOT / 'weights' / xdb).exists() and t2 < float(mcfg.get('extra_guard', 330)): try: extra = run_diamond(write_chunks(enumerate(seqs), workdir, 'x'), mcfg.get('extra_mode', 'fast'), cfg, workdir, db=xdb) for i, h in extra.items(): hits.setdefault(i, []).extend(h) log(f'extra reference pass ({xdb}) done: {len(extra)} inputs with hits, {time.time() - START - t2:.0f}s') except Exception as exc: log(f'extra reference pass skipped: {exc}') elif xdb: log(f'extra reference pass skipped (elapsed {t2:.0f}s)') finally: shutil.rmtree(workdir, ignore_errors=True) # Scoring mode (weights/mode.json): "learned" = gradient-boosted re-ranker (default), # "vote" = plain bitscore^power vote over the top-k hits. mode_cfg = load_mode() log(f'scoring mode: {mode_cfg}') if mode_cfg['scoring'] == 'vote': preds = {} for i in sorted(hits): h = sorted(hits[i], key=lambda x: -x[5])[:int(mode_cfg['vote_k'])] votes = {} srcw = mode_cfg.get('source_weights', {}) # e.g. down-weight unreviewed (TrEMBL) references for s, pid, ql, sl, ev, b, *_ in h: ecs = labels[s] w = (b / h[0][5]) ** float(mode_cfg['vote_power']) * srcw.get(s[:3], 1.0) for e in ecs: votes[e] = votes.get(e, 0.0) + w / len(ecs) if votes: preds[i] = sorted(votes, key=lambda e: (-votes[e], -freq.get(e, 0), e))[:10] # v17: learned re-ranker over the vote's top 10 candidates, used only when the best hit is below # rerank_max_pid identity (it was trained and validated on official-label hard proteins there). rr_path = ROOT / 'weights' / 'ranker_v2.npz' if mode_cfg.get('reranker') and rr_path.exists(): try: from rerank2 import candidate_features rr = TreeRanker(rr_path) srcw = mode_cfg.get('source_weights', {}) qs, blocks, cand_of = [], [], {} for i in sorted(hits): if not hits[i] or max(x[1] for x in hits[i]) >= float(mode_cfg.get('rerank_max_pid', 80)): continue cands, X = candidate_features(hits[i], labels, freq, srcw, int(mode_cfg['vote_k']), float(mode_cfg['vote_power'])) if len(cands) > 1: qs.append(i) cand_of[i] = cands blocks.append(np.asarray(X, dtype=np.float64)) if blocks: scores = rr.score(np.vstack(blocks)) pos = 0 for i in qs: c = cand_of[i] sc = scores[pos:pos + len(c)] pos += len(c) preds[i] = [c[j] for j in np.argsort(-sc, kind='stable')] log(f're-ranker v2 applied to {len(qs)} inputs') except Exception as exc: log(f're-ranker v2 skipped, using the vote: {exc!r}') # v24: exact-sequence label table (weights/srchash.npz, built by tools/build_source_hash.py). For an input with no # identical train / Swiss-Prot sequence, the table's label (source priority measured on the training labels: # BRENDA > CatPred-DB > CataPro > DLKcat > PDB/SIFTS > UniProtKB unreviewed > ESP) is placed first. hp = ROOT / 'weights' / str(mode_cfg.get('label_table', 'srchash.npz')) if mode_cfg.get('label_table') and hp.exists(): try: import hashlib z = np.load(hp) H, S = z['h'], z['s'] sets = json.loads(hp.with_name(hp.stem + '_sets.json').read_text()) used = 0 for i, sq in enumerate(seqs): if any(x[1] >= 100.0 and x[2] == x[3] and x[7] - x[6] + 1 == x[2] and not x[0].startswith(('un_', 'u2_', 'cp_', 'br_', 'dk_', 'es_', 'e2_', 'bo_', 'pd_')) for x in hits.get(i, [])): continue # identical train / Swiss-Prot sequence: keep its label key = np.uint64(int(hashlib.md5(sq.encode()).hexdigest()[:16], 16)) j = int(np.searchsorted(H, key)) if j < len(H) and H[j] == key: ecs = sets[int(S[j])] rank = {e: r for r, e in enumerate(preds.get(i, []))} ordered = sorted(ecs, key=lambda e: (rank.get(e, 99), -freq.get(e, 0), e)) preds[i] = (ordered + [e for e in preds.get(i, []) if e not in ecs])[:10] used += 1 log(f'exact-sequence label table: {len(H)} sequences, used for {used} inputs') except Exception as exc: log(f'exact-sequence label table skipped: {exc!r}') return preds ranker = TreeRanker(ROOT / 'weights' / 'ranker.npz') rel = EcReliability(ROOT / 'weights' / 'te_tables.json.gz') order, cand_of, blocks = sorted(hits), {}, [] for i in order: cands, X = query_candidates(sorted(hits[i], key=lambda x: -x[5]), labels, freq) if cands: cand_of[i] = cands blocks.append(np.hstack([X.astype(np.float64), rel.matrix(cands)])) scores = ranker.score(np.vstack(blocks)) if blocks else np.zeros(0) # v7: add CNN features and use the CNN-aware re-ranker when available and time allows. cnn_w, cnn_r = ROOT / 'weights' / 'cnn.npz', ROOT / 'weights' / 'ranker_cnn.npz' if blocks and cnn_w.exists() and cnn_r.exists() and time.time() - START < 450: try: from cnn_numpy import SmallCNN from cnn_feats import CnnScorer t0 = time.time() net = SmallCNN(cnn_w) scorer = CnnScorer(net.classes) qs = [i for i in order if i in cand_of] L = net.logits([seqs[i] for i in qs]) blocks2 = [np.hstack([blk, scorer.rows(L[k], cand_of[i])]) for k, (i, blk) in enumerate(zip(qs, blocks))] scores = TreeRanker(cnn_r).score(np.vstack(blocks2)) log(f'CNN re-ranker used ({time.time() - t0:.0f}s)') except Exception as exc: log(f'CNN step skipped, using v5 re-ranker: {exc!r}') preds, pos = {}, 0 for i in order: if i not in cand_of: continue cands = cand_of[i] s = scores[pos:pos + len(cands)] pos += len(cands) preds[i] = [cands[j] for j in np.argsort(-s, kind='stable')] return preds def kmer_predictions(seqs, labels, freq): """Fallback without DIAMOND: 3-mer TF-IDF nearest neighbours, bitscore-free vote.""" from kmer_knn import KmerIndex ids, train_seqs = [], [] for part in sorted((ROOT / 'weights').glob('train_seqs_*.tsv.gz')): with gzip.open(part, 'rt') as f: for line in f: i, s = line.rstrip('\n').split('\t') ids.append(i) train_seqs.append(s) log(f'k-mer fallback: {len(train_seqs)} training sequences') idx = KmerIndex(train_seqs) del idx.XT, train_seqs I, S = idx.query_fast(seqs) preds = {} for q, (row, sc) in enumerate(zip(I, S)): votes = {} for j, s in zip(row, sc): ecs = labels[ids[j]] for e in ecs: votes[e] = votes.get(e, 0.0) + max(float(s), 0.0) ** 4 / len(ecs) if votes: preds[q] = sorted(votes, key=lambda e: (-votes[e], -freq.get(e, 0), e))[:10] return preds def predict_file(input_path, output_path): diagnostics() cpu_probe() frame = prepare_inputs(pd.read_parquet(input_path)) with gzip.open(ROOT / 'weights' / 'labels.json.gz', 'rt') as f: model = json.load(f) labels, freq = model['labels'], model['freq'] default = sorted(freq, key=lambda e: (-freq[e], e))[0] seqs = list(frame.sequence) log(f'{len(seqs)} unique inputs') preds, method = None, 'diamond' try: preds = homology_predictions(seqs, labels, freq) except Exception as exc: log(f'homology pipeline error: {exc!r}') if preds is None: method = 'kmer-fallback' log('DIAMOND unavailable -> k-mer fallback') try: preds = kmer_predictions(seqs, labels, freq) except Exception as exc: log(f'k-mer fallback error: {exc!r}') preds, method = {}, 'frequency-only' with Path(output_path).open('w') as out: for i, key in enumerate(frame.id): out.write(json.dumps({'id': key, 'prediction': preds.get(i) or [default]}) + '\n') log(f'done: method={method}, {len(preds)}/{len(seqs)} predicted from neighbours') 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)