From e95bdeafea917348c7fee7b971990e94546116ba Mon Sep 17 00:00:00 2001 From: wdecoster Date: Sun, 14 Jun 2026 10:43:22 +0200 Subject: [PATCH] feat: per-locus CPU-time profiling (thread CPU time + harness) Measure per-locus genotyping cost by thread CPU time instead of wall clock, so timings are not inflated when other processes contend for the machine. Each locus is genotyped start-to-finish on a single rayon worker, so the thread-CPU delta around it is that locus's CPU cost. - repeats.rs: add thread_cpu_time() (clock_gettime(CLOCK_THREAD_CPUTIME_ID), hence the libc dep), store it in RepeatInterval.created (now a Duration), and add cpu_elapsed()/time_field() helpers. - vcf.rs: the debug-only ;TIME= INFO field now reports thread CPU seconds; the four duplicated timing blocks collapse to repeat.time_field(args.debug). - misc/profile_loci.py: run STRdust over N replicates, drop warmup runs, aggregate per-locus median/MAD/CV CPU time and correlate with VCF features (depth, allele length, stdev, cluster count, ...) to find consistently slow loci. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 1 + Cargo.toml | 1 + misc/profile_loci.py | 281 +++++++++++++++++++++++++++++++++++++++++++ src/repeats.rs | 39 +++++- src/vcf.rs | 42 +------ 5 files changed, 325 insertions(+), 39 deletions(-) create mode 100644 misc/profile_loci.py diff --git a/Cargo.lock b/Cargo.lock index 80c98cf..036a17d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14,6 +14,7 @@ dependencies = [ "hts-sys", "indicatif", "kodama", + "libc", "log", "minimap2", "rand 0.10.1", diff --git a/Cargo.toml b/Cargo.toml index 3f8ecfa..a87eecf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,7 @@ rand = "0.10.1" hts-sys = "2.1.1" reqwest = { version = "0.13.2", features = ["blocking", "json"] } indicatif = { version = "0.18.0", features = ["rayon"] } +libc = "0.2" [dev-dependencies] diff --git a/misc/profile_loci.py b/misc/profile_loci.py new file mode 100644 index 0000000..141d205 --- /dev/null +++ b/misc/profile_loci.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 +"""Profile STRdust per-locus CPU cost across replicates and relate it to locus features. + +STRdust emits a per-locus ``TIME=s`` INFO field in ``--debug`` mode. Since +commit "perf: thread CPU time" that value is *thread CPU time* (CLOCK_THREAD_CPUTIME_ID), +not wall-clock, so it is largely immune to other processes stealing the core. Even so a +single run is noisy, so this script runs STRdust several times and aggregates per locus to +find the ones that are *consistently* slow, then correlates slowness with locus features +parsed from the same VCF (depth, allele length, stdev, cluster count, ...). + +Typical use (run on the machine with the data, not the laptop): + + python misc/profile_loci.py \ + --binary ./target/release/STRdust \ + --fasta ref.fa --bam sample.cram --bed loci.bed \ + --replicates 5 --threads 1 \ + --out-prefix profile_run1 + +Outputs: + .per_locus.tsv one row per locus: median/MAD/CV CPU time + features + .replicates.tsv raw per-locus time for every replicate (long format) + .png (optional, --plot) feature-vs-time scatter panels + +Only the standard library is required to *run the replicates and parse*. pandas + scipy are +used for the aggregation/correlation report and are imported lazily, so a parse-only run +(``--from-vcfs``) still works in a bare environment. +""" + +import argparse +import subprocess +import sys +import re +import statistics +from pathlib import Path + +TIME_RE = re.compile(r"TIME=([0-9.]+)s") + + +def parse_info(info: str) -> dict: + """Parse a VCF INFO column into a dict; flags map to True.""" + out = {} + for field in info.split(";"): + if not field: + continue + if "=" in field: + k, v = field.split("=", 1) + out[k] = v + else: + out[field] = True + return out + + +def locus_features(line: str) -> dict | None: + """Extract time + features from one VCF data line, or None if it has no TIME field.""" + cols = line.rstrip("\n").split("\t") + if len(cols) < 8: + return None + chrom, pos, _id, ref, alt, _qual, _filt, info = cols[:8] + m = TIME_RE.search(info) + if m is None: + return None # not a --debug line / no timing + info_d = parse_info(info) + + # FORMAT/sample (present for genotyped loci, absent for some missing records) + fmt = cols[8].split(":") if len(cols) > 8 else [] + sample = cols[9].split(":") if len(cols) > 9 else [] + sample_d = dict(zip(fmt, sample)) + + def two_max(field, cast=int): + """Max of a 2-value 'a,b' FORMAT field, ignoring '.'.""" + vals = [cast(x) for x in sample_d.get(field, "").split(",") if x not in (".", "")] + return max(vals) if vals else None + + def two_sum(field, cast=int): + vals = [cast(x) for x in sample_d.get(field, "").split(",") if x not in (".", "")] + return sum(vals) if vals else None + + alt_alleles = [a for a in alt.split(",") if a != "."] + outliers = info_d.get("OUTLIERS") + + return { + "locus": f"{chrom}:{pos}", + "chrom": chrom, + "pos": int(pos), + "time_s": float(m.group(1)), + "ref_len": len(ref) if ref != "." else None, + "n_alt": len(alt_alleles), + "max_alt_len": max((len(a) for a in alt_alleles), default=0), + # RB = length relative to ref, FRB = full repeat length, SUP = read support + "max_rb": two_max("RB"), + "max_frb": two_max("FRB"), + "total_sup": two_sum("SUP"), + "max_stdev": max((int(x) for x in info_d.get("STDEV", "").split(",") + if x not in (".", "")), default=None), + "nclusters": int(info_d["NCLUSTERS"]) if "NCLUSTERS" in info_d else None, + "n_outliers": len(outliers.split(",")) if outliers else 0, + "quickref": "QUICKREF" in info_d, + "clusterfailure": "CLUSTERFAILURE" in info_d, + } + + +def run_replicate(binary, fasta, bam, bed, region, threads, debug_extra): + """Run STRdust once with --debug and return its VCF stdout as a list of lines.""" + cmd = [binary, "--debug", "-t", str(threads)] + if bed: + cmd += ["-R", bed] + if region: + cmd += ["-r", region] + cmd += list(debug_extra) + [fasta, bam] + print(f" $ {' '.join(cmd)}", file=sys.stderr) + proc = subprocess.run(cmd, capture_output=True, text=True) + if proc.returncode != 0: + sys.exit(f"STRdust failed (exit {proc.returncode}):\n{proc.stderr[-2000:]}") + return proc.stdout.splitlines() + + +def collect(lines): + """Map locus -> features for every timed data line in one VCF.""" + out = {} + for line in lines: + if line.startswith("#"): + continue + feat = locus_features(line) + if feat is not None: + out[feat["locus"]] = feat + return out + + +def aggregate(replicates, warmup, out_prefix): + """Aggregate per-locus times across replicates and write TSV reports.""" + # Drop warmup replicates (cold cache) from the timing aggregation. + timed = replicates[warmup:] + if not timed: + sys.exit("No replicates left after dropping warmup; lower --warmup.") + + # Raw long-format dump: locus, replicate index, time. + long_path = Path(f"{out_prefix}.replicates.tsv") + with long_path.open("w") as fh: + fh.write("locus\treplicate\ttime_s\n") + for i, rep in enumerate(replicates): + tag = "warmup" if i < warmup else str(i - warmup) + for locus, feat in rep.items(): + fh.write(f"{locus}\t{tag}\t{feat['time_s']:.6f}\n") + + # Features come from the last non-warmup replicate (genotypes are deterministic). + feature_src = timed[-1] + loci = sorted(feature_src, key=lambda k: (feature_src[k]["chrom"], feature_src[k]["pos"])) + + rows = [] + for locus in loci: + times = [rep[locus]["time_s"] for rep in timed if locus in rep] + if not times: + continue + med = statistics.median(times) + mad = statistics.median([abs(t - med) for t in times]) if len(times) > 1 else 0.0 + cv = (statistics.pstdev(times) / med) if len(times) > 1 and med > 0 else 0.0 + f = feature_src[locus] + rows.append({ + "locus": locus, "chrom": f["chrom"], "pos": f["pos"], + "n_rep": len(times), + "median_s": med, "min_s": min(times), "max_s": max(times), + "mad_s": mad, "cv": cv, + "ref_len": f["ref_len"], "max_rb": f["max_rb"], "max_frb": f["max_frb"], + "max_alt_len": f["max_alt_len"], "total_sup": f["total_sup"], + "max_stdev": f["max_stdev"], "nclusters": f["nclusters"], + "n_outliers": f["n_outliers"], "n_alt": f["n_alt"], + "quickref": int(f["quickref"]), "clusterfailure": int(f["clusterfailure"]), + }) + + cols = list(rows[0].keys()) + per_locus_path = Path(f"{out_prefix}.per_locus.tsv") + with per_locus_path.open("w") as fh: + fh.write("\t".join(cols) + "\n") + for r in rows: + fh.write("\t".join("" if r[c] is None else + (f"{r[c]:.6f}" if isinstance(r[c], float) else str(r[c])) + for c in cols) + "\n") + + print(f"\nWrote {per_locus_path} ({len(rows)} loci) and {long_path}", file=sys.stderr) + return rows, per_locus_path + + +def report(rows, top, per_locus_path, plot, out_prefix): + """Print the consistently-slowest loci and feature correlations; optional scatter plot.""" + consistent = sorted(rows, key=lambda r: r["median_s"], reverse=True) + print(f"\n=== Top {top} consistently slowest loci (by median thread-CPU time) ===") + hdr = ("locus", "median_s", "cv", "total_sup", "max_frb", "nclusters", "n_outliers") + print("\t".join(hdr)) + for r in consistent[:top]: + print("\t".join(str(r[c] if r[c] is not None else "") for c in hdr)) + + # Correlations need scipy/pandas; degrade gracefully if absent. + try: + import pandas as pd + from scipy.stats import spearmanr + except ImportError: + print("\n(install pandas + scipy for the correlation report)", file=sys.stderr) + return + + df = pd.read_csv(per_locus_path, sep="\t") + numeric = ["ref_len", "max_rb", "max_frb", "max_alt_len", "total_sup", + "max_stdev", "nclusters", "n_outliers", "n_alt", "quickref"] + print("\n=== Spearman correlation of features with median CPU time ===") + print("feature\trho\tp_value\tn") + corrs = [] + for col in numeric: + sub = df[[col, "median_s"]].dropna() + if sub[col].nunique() < 2 or len(sub) < 3: + continue + rho, p = spearmanr(sub[col], sub["median_s"]) + corrs.append((col, rho, p, len(sub))) + for col, rho, p, n in sorted(corrs, key=lambda x: abs(x[1]), reverse=True): + print(f"{col}\t{rho:+.3f}\t{p:.2e}\t{n}") + + if plot: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + feats = [c for c, *_ in sorted(corrs, key=lambda x: abs(x[1]), reverse=True)][:6] + n = len(feats) + if n: + ncol = 3 + nrow = (n + ncol - 1) // ncol + fig, axes = plt.subplots(nrow, ncol, figsize=(4 * ncol, 3.2 * nrow), squeeze=False) + for ax, col in zip(axes.flat, feats): + ax.scatter(df[col], df["median_s"], s=8, alpha=0.4) + ax.set_xlabel(col) + ax.set_ylabel("median CPU time (s)") + for ax in axes.flat[n:]: + ax.set_visible(False) + fig.tight_layout() + png = f"{out_prefix}.png" + fig.savefig(png, dpi=130) + print(f"\nWrote {png}", file=sys.stderr) + + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--binary", default="./target/release/STRdust", help="STRdust binary") + p.add_argument("--fasta", help="reference fasta") + p.add_argument("--bam", help="BAM/CRAM file") + p.add_argument("--bed", help="region bed file (-R)") + p.add_argument("--region", help="single region chr:start-end (-r)") + p.add_argument("--replicates", type=int, default=5, help="number of STRdust runs") + p.add_argument("--warmup", type=int, default=1, + help="leading replicates to exclude from timing (cold cache)") + p.add_argument("--threads", type=int, default=1, + help="threads per run; 1 gives the least cache contention") + p.add_argument("--debug-extra", nargs=argparse.REMAINDER, default=[], + help="extra args passed verbatim before FASTA/BAM (e.g. --phasing-strategy dbscan)") + p.add_argument("--from-vcfs", nargs="+", + help="skip running; aggregate these already-produced --debug VCFs instead") + p.add_argument("--out-prefix", default="strdust_profile", help="output file prefix") + p.add_argument("--top", type=int, default=25, help="how many slow loci to print") + p.add_argument("--plot", action="store_true", help="write a feature-vs-time scatter PNG") + args = p.parse_args() + + if args.from_vcfs: + replicates = [collect(Path(v).read_text().splitlines()) for v in args.from_vcfs] + print(f"Parsed {len(replicates)} VCF(s)", file=sys.stderr) + else: + if not (args.fasta and args.bam and (args.bed or args.region)): + p.error("need --fasta, --bam and one of --bed/--region (or use --from-vcfs)") + if args.threads != 1: + print(f"warning: --threads {args.threads}: per-thread CPU time stays valid, " + "but cache contention between workers adds noise; --threads 1 is cleanest.", + file=sys.stderr) + replicates = [] + for i in range(args.replicates): + print(f"replicate {i + 1}/{args.replicates}", file=sys.stderr) + lines = run_replicate(args.binary, args.fasta, args.bam, args.bed, + args.region, args.threads, args.debug_extra) + replicates.append(collect(lines)) + + rows, per_locus_path = aggregate(replicates, args.warmup, args.out_prefix) + report(rows, args.top, per_locus_path, args.plot, args.out_prefix) + + +if __name__ == "__main__": + main() diff --git a/src/repeats.rs b/src/repeats.rs index 549c3b8..821744c 100644 --- a/src/repeats.rs +++ b/src/repeats.rs @@ -245,7 +245,26 @@ pub struct RepeatInterval { pub chrom: String, pub start: u32, pub end: u32, - pub created: Option, + /// Thread CPU time consumed at the moment genotyping of this locus started + /// (see [`thread_cpu_time`]). Subtracting it from a later reading on the same + /// thread yields the locus's CPU cost, reported via the `TIME=` INFO field. + pub created: Option, +} + +/// Read the calling thread's consumed CPU time. +/// +/// Unlike `Instant::now()` (wall-clock), this clock only advances while the +/// thread is actually scheduled on a CPU, so per-locus measurements are not +/// inflated by other processes competing for the machine. Each locus is +/// genotyped start-to-finish on a single rayon worker thread, so the delta +/// between two readings on that thread is the locus's CPU cost. +pub fn thread_cpu_time() -> std::time::Duration { + let mut ts = libc::timespec { tv_sec: 0, tv_nsec: 0 }; + // SAFETY: `ts` is a valid, fully-initialized timespec; clock_gettime only + // writes the seconds/nanoseconds fields and returns 0 on success. + let rc = unsafe { libc::clock_gettime(libc::CLOCK_THREAD_CPUTIME_ID, &mut ts) }; + debug_assert_eq!(rc, 0, "clock_gettime(CLOCK_THREAD_CPUTIME_ID) failed"); + std::time::Duration::new(ts.tv_sec as u64, ts.tv_nsec as u32) } impl fmt::Display for RepeatInterval { @@ -353,7 +372,23 @@ impl RepeatInterval { } pub fn set_time_stamp(&mut self) { - self.created = Some(std::time::Instant::now()); + self.created = Some(thread_cpu_time()); + } + + /// CPU time consumed since [`set_time_stamp`](Self::set_time_stamp) was + /// called, or `None` if it never was. + pub fn cpu_elapsed(&self) -> Option { + self.created + .map(|start| thread_cpu_time().saturating_sub(start)) + } + + /// Render the `;TIME=s` INFO field. Empty unless `debug` is set + /// and a start timestamp was recorded. + pub fn time_field(&self, debug: bool) -> String { + match (debug, self.cpu_elapsed()) { + (true, Some(elapsed)) => format!(";TIME={:.3}s", elapsed.as_secs_f64()), + _ => String::new(), + } } } diff --git a/src/vcf.rs b/src/vcf.rs index 8811a00..76c7fd3 100644 --- a/src/vcf.rs +++ b/src/vcf.rs @@ -131,15 +131,7 @@ impl VCFRecord { } else { format!("{};", flag.join(";")) }; - let time_taken = if args.debug { - let elapsed = repeat - .created - .expect("Failed accessing timestamp") - .elapsed(); - format!(";TIME={:.3}s", elapsed.as_secs_f64()) - } else { - "".to_string() - }; + let time_taken = repeat.time_field(args.debug); VCFRecord { chrom: repeat.chrom.clone(), start: repeat.start, @@ -166,15 +158,7 @@ impl VCFRecord { flags: Vec, args: &crate::Cli, ) -> VCFRecord { - let time_taken = if args.debug { - let elapsed = repeat - .created - .expect("Failed accessing timestamp") - .elapsed(); - format!(";TIME={:.3}s", elapsed.as_secs_f64()) - } else { - "".to_string() - }; + let time_taken = repeat.time_field(args.debug); let flags_str = if flags.is_empty() { "".to_string() } else { @@ -206,15 +190,7 @@ impl VCFRecord { support: String, args: &crate::Cli, ) -> VCFRecord { - let time_taken = if args.debug { - let elapsed = repeat - .created - .expect("Failed accessing timestamp") - .elapsed(); - format!(";TIME={:.3}s", elapsed.as_secs_f64()) - } else { - "".to_string() - }; + let time_taken = repeat.time_field(args.debug); VCFRecord { chrom: repeat.chrom.clone(), start: repeat.start, @@ -250,15 +226,7 @@ impl VCFRecord { } None => "".to_string(), }; - let time_taken = if args.debug { - let elapsed = repeat - .created - .expect("Failed accessing timestamp") - .elapsed(); - format!(";TIME={:.3}s", elapsed.as_secs_f64()) - } else { - "".to_string() - }; + let time_taken = repeat.time_field(args.debug); VCFRecord { chrom: repeat.chrom.clone(), start: repeat.start, @@ -538,7 +506,7 @@ pub fn write_vcf_header(args: &Cli) { ); if args.debug { println!( - r#"##INFO="# + r#"##INFO="# ); } println!(r#"##FORMAT="#);