diff --git a/metagraph/docs/PREFETCH_AGENT_HANDOFF.md b/metagraph/docs/PREFETCH_AGENT_HANDOFF.md new file mode 100644 index 0000000000..851ee6dacc --- /dev/null +++ b/metagraph/docs/PREFETCH_AGENT_HANDOFF.md @@ -0,0 +1,120 @@ +# Agent handoff: mmap prefetch work (`mk/prefetch`) + +This note is for whoever continues **MADV_WILLNEED prefetch** work and runs benchmarks on a server. Read this first, then open the PR linked below for the full checklist and **Results** table. + +## TL;DR + +| Item | Value | +|------|--------| +| **Branch** | `mk/prefetch` (starts from `origin/master`) | +| **PR** | [#628](https://github.com/ratschlab/metagraph/pull/628) — implementation plan + phase checkboxes + **Results** table | +| **Code landed (local)** | Phases **0–5** on `mk/prefetch` (see commits below). Phases **6–7** still open per PR. | +| **Next code step** | **Phase 6** (optional `valid_edges_`, data-driven) **or** **Phase 7** (virtual `prefetch()` cleanup on `DeBruijnGraph` / `BinaryMatrix` after more phases merge). | + +### Commits on `mk/prefetch` (prefetch stack) + +| Phase | Commit | Summary | +|-------|--------|---------| +| **0** | `532ea6de` | `scripts/bench_query_prefetch.py` | +| **0** (doc) | `701d67c7` | This handoff + pointers | +| **1** | `f4cdb48c` | `suffix_ranges` + `utils::madvise_willneed` / `get_mmap_data` + `query_fasta` (CanonicalDBG unwrap) | +| **2** | `9443f9de` | RowDiff `.anchors` / `.rd_succ` mmap + `IRowDiff::prefetch()` + `query_fasta` | +| **3–5** | `e2a8c656` | Bloom `.bloom` mmap + `DBGSuccinct::prefetch_bloom_filter`; BRWT `prefetch_if_dense`; RowDisk / IntRowDisk / CoordRowDisk boundary `MADV_WILLNEED` at batch entry | + +## What you should do first (server) + +1. **Checkout and build** + ```bash + git fetch origin && git checkout mk/prefetch && git pull + # build metagraph as usual (e.g. cmake + make in your build dir) + ``` + Use the DNA binary for benchmarks, e.g. `.../build/metagraph_DNA`. + +2. **Baseline benchmark (Phase 0 row in PR table)** — compare against **`701d67c7`** (script + docs, **no** C++ prefetch) **or** rebuild that commit for a fair binary match. + ```bash + metagraph/scripts/bench_query_prefetch.py \ + --metagraph /path/to/your/build/metagraph_DNA \ + --graphs graphs.csv \ + --query reads.fa \ + --mmap --madv-random \ + -p 4 --threads-each 8 \ + --warmup 1 --repeats 3 \ + --json bench/phase0_baseline.json + ``` + Optional **cold OS page cache** between repeats (needs root): + ```bash + sudo metagraph/scripts/bench_query_prefetch.py ... --drop-cache --json bench/phase0_cold.json + ``` + **Note:** `sudo` is **only** for `--drop-cache`. Normal runs do not need root. + +3. **Record results** in PR #628 → **Results** table (one row per phase you benchmark). + +## Benchmark script + +- **Path:** `metagraph/scripts/bench_query_prefetch.py` +- **Behavior:** Spawns `metagraph server_query`, POSTs `/search` with FASTA JSON (same shape as `api/python/metagraph/client.py`). Parses `-v` trace lines: + - **K-mer mapping:** `[Query graph construction] Contigs mapped to the full graph ... in X sec` + - **Row-diff annotation:** `RD query [...] traversal: X sec, call_rd_rows: Y sec, ...` +- **Output:** Human-readable summary + `--json` for diffing runs. +- **Docs:** Module docstring at top of the script. + +### Multi-graph CSV (`server_query`) + +One row per index: + +``` +,, +``` + +Spaces around commas break parsing — keep rows tight. See `server_query` help in `src/cli/config/config.cpp`. + +**`graphs` JSON filter:** If the CSV has **more than 10 distinct names**, the server rejects requests **without** a `"graphs"` list. The script’s repeated **`--graph-name`** flags supply that list (one flag per distinct name you want to query). + +### Flags that matter for prefetch A/B + +- **`--mmap`** — graph/annotation loaded via mmap (prefetch targets exist). +- **`--madv-random`** — enables madvise hints (`utils::with_madvise()`); all `MADV_WILLNEED` paths are gated the same way. + +Forward extra server flags with repeated `--server-arg`, e.g.: + +```bash +--server-arg --query-batch-size --server-arg 100000 +``` + +## Implementation roadmap (sync with PR #628) + +The **authoritative** checklist is in **PR #628**. Status on this branch: + +| Phase | Content | Status on `mk/prefetch` | +|-------|---------|-------------------------| +| **0** | Benchmark script | Done (`532ea6de`) | +| **1** | `suffix_ranges` + `madvise_willneed` + `get_mmap_data` + CanonicalDBG unwrap | Done (`f4cdb48c`; originally cherry-picked from `94d7caa60` / `mk/madvise-suffix-ranges`) | +| **2** | RowDiff `anchor_` / `fork_succ_` prefetch | Done (`9443f9de`) | +| **3** | Bloom filter prefetch | Done (`e2a8c656`) | +| **4** | BRWT adaptive `prefetch_if_dense` on `nonzero_rows_` | Done (`e2a8c656`) | +| **5** | RowDisk / IntRowDisk / CoordRowDisk `boundary_` prefetch | Done (`e2a8c656`) | +| **6** | Optional `valid_edges_` | **Not started** (data-driven) | +| **7** | Virtual `prefetch()` cleanup | **Not started** (after several phases land) | + +After **each** phase you benchmark: rebuild at that commit, re-run the script with the **same** arguments, save a new `--json`, update the PR **Results** table and tick the checkbox in PR #628. + +## Gotchas + +1. **Trace level:** The script always passes **`-v`** to `server_query` so trace lines appear. +2. **Per-run attribution:** The script slices server log by buffer offsets between timed requests. **Overlapping** concurrent clients can blur per-run stats; single-client `--parallel 1` on the script side is safest. +3. **`--drop-cache`:** Without root, the script warns and skips cache drop; benchmark still runs. For A/B on warm servers, skipping cache drop is often fine (see PR discussion). +4. **macOS vs Linux build dirs:** Adjust `--metagraph` to your `metagraph_DNA` path. +5. **Single-graph `server_query` + scripts:** The server may log **“Will listen”** before the async graph load finishes; readiness is **HTTP `GET /stats` returning 200** (not only a log grep). Multi-graph mode logs **“Ready to serve queries”** after load. +6. **`strace -e madvise`:** Shows `MADV_WILLNEED` when prefetches run; most wall time is still normal I/O/CPU outside `madvise`. + +## Related code pointers + +- Server multi-graph load: `src/cli/server.cpp` (CSV parsing, `graphs_cache`) +- Search endpoint: `POST /search` in same file +- **Phase 1:** `src/common/utils/file_utils.{hpp,cpp}`, `src/graph/representation/succinct/dbg_succinct.{hpp,cpp}`, `src/cli/query.cpp` +- **Phase 2:** `src/annotation/binary_matrix/row_diff/row_diff.{hpp,cpp}`, `src/cli/query.cpp` +- **Phases 3–5:** `src/kmer/kmer_bloom_filter.{hpp,cpp}`, `src/graph/representation/succinct/dbg_succinct.{hpp,cpp}`, `src/annotation/binary_matrix/multi_brwt/brwt.{hpp,cpp}`, `src/annotation/binary_matrix/row_disk/row_disk.{hpp,cpp}`, `src/annotation/int_matrix/row_disk/{int_row_disk,coord_row_disk}.{hpp,cpp}`, `src/cli/query.cpp` + +## Questions? + +If PR #628 body and this file disagree, **prefer PR #628** for checklist/results — update this doc when the workflow changes. diff --git a/metagraph/scripts/bench_query_prefetch.py b/metagraph/scripts/bench_query_prefetch.py new file mode 100755 index 0000000000..df8b0f72bd --- /dev/null +++ b/metagraph/scripts/bench_query_prefetch.py @@ -0,0 +1,716 @@ +#!/usr/bin/env python3 +"""Benchmark `metagraph server_query` query latency, broken down per stage. + +Spawns a `metagraph server_query` process against a multi-graph CSV (or a +single -i/-a pair), sends one or more `/search` requests over HTTP, and +parses the server's trace logs to summarize: + + * "[Query graph construction] Contigs mapped to the full graph ..." + -> k-mer mapping latency per batch + * "RD query [...] -- traversal: ... call_rd_rows: ..." + -> annotation (row-diff) traversal + decoding latency per batch + +Designed for A/B testing optimizations such as MADV_WILLNEED prefetch: +run the same workload before and after a change, optionally with the OS +page cache dropped between runs, and diff the resulting summaries (or +the `--json` output). + +Stdlib only -- no `requests` / `pandas` dependency. + +Example: + + scripts/bench_query_prefetch.py \\ + --metagraph build/metagraph/metagraph \\ + --graphs graphs.csv \\ + --query reads.fa \\ + --mmap --madv-random \\ + -p 4 --threads-each 8 \\ + --warmup 1 --repeats 3 \\ + --json results_after.json + +Add `--drop-cache` to evict the OS page cache between repeats (requires +root: writes `/proc/sys/vm/drop_caches` on Linux, runs `purge` on macOS). +Without root the flag is a no-op with a warning -- the rest of the script +runs unprivileged. +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import re +import shlex +import signal +import statistics +import subprocess +import sys +import tempfile +import threading +import time +import urllib.error +import urllib.request +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + + +# --------------------------------------------------------------------------- +# Trace-line regexes. Anchored to the substantive part of the log line so +# changes to the "[YYYY-MM-DD HH:MM:SS.mmm] [trace] " prefix don't break us. +# --------------------------------------------------------------------------- + +KMER_MAP_RE = re.compile( + r"\[Query graph construction\] Contigs mapped to the full graph " + r"\[threads: (?P\d+), contigs: (?P\d+), " + r"chunk_size: (?P\d+)\] " + r"\(found (?P\d+) / (?P\d+) k-mers\) " + r"in (?P[\d.]+) sec" +) + +RD_QUERY_RE = re.compile( + r"RD query \[threads: (?P\d+), " + r"rows: (?P\d+) -> (?P\d+) " + r"\((?P[\d.]+)x\)\] -- " + r"traversal: (?P[\d.]+) sec, " + r"call_rd_rows: (?P[\d.]+) sec " + r"\(set bits: (?P\d+), capacity: (?P\d+)\), " + r"decoding: (?P[\d.]+) sec, " + r"reconstruction: (?P[\d.]+) sec" +) + +# server.cpp emits this once all graphs have been loaded. +SERVER_READY_RE = re.compile(r"All graphs were loaded.*Ready to serve queries") + +# server.cpp also emits this single-graph variant when -i/-a is used (no CSV). +SINGLE_GRAPH_READY_RE = re.compile(r"\[Server\] Will listen on") + + +# --------------------------------------------------------------------------- +# Stats helpers +# --------------------------------------------------------------------------- + + +def summarize(xs: List[float]) -> Dict[str, float]: + """Return count / mean / median / p50 / p90 / p99 / max / total of `xs`.""" + if not xs: + return {"count": 0} + s = sorted(xs) + return { + "count": len(xs), + "mean": statistics.fmean(xs), + "median": statistics.median(s), + "p50": _percentile(s, 50), + "p90": _percentile(s, 90), + "p99": _percentile(s, 99), + "max": s[-1], + "total": sum(xs), + } + + +def _percentile(sorted_xs: List[float], pct: float) -> float: + if not sorted_xs: + return 0.0 + if len(sorted_xs) == 1: + return sorted_xs[0] + # linear interpolation between closest ranks (numpy's default) + k = (len(sorted_xs) - 1) * (pct / 100.0) + lo = int(k) + hi = min(lo + 1, len(sorted_xs) - 1) + frac = k - lo + return sorted_xs[lo] * (1 - frac) + sorted_xs[hi] * frac + + +def fmt_summary(label: str, units: str, s: Dict[str, float]) -> str: + if not s.get("count"): + return f" {label:24s} (no events)" + return ( + f" {label:24s} " + f"n={s['count']:<5d} " + f"mean={s['mean']:.3f}{units} " + f"med={s['median']:.3f}{units} " + f"p99={s['p99']:.3f}{units} " + f"max={s['max']:.3f}{units} " + f"total={s['total']:.2f}{units}" + ) + + +# --------------------------------------------------------------------------- +# Server lifecycle +# --------------------------------------------------------------------------- + + +@dataclass +class ServerProcess: + proc: subprocess.Popen + log_path: Path + drain_thread: threading.Thread + log_buffer: List[str] = field(default_factory=list) + log_lock: threading.Lock = field(default_factory=threading.Lock) + + def snapshot(self) -> str: + """Return everything captured from the server's stderr so far.""" + with self.log_lock: + return "".join(self.log_buffer) + + def stop(self, timeout: float = 10.0) -> None: + if self.proc.poll() is None: + self.proc.send_signal(signal.SIGTERM) + try: + self.proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + self.proc.kill() + self.proc.wait() + self.drain_thread.join(timeout=2.0) + + +def spawn_server( + metagraph: str, + graphs_csv: Optional[str], + single_index: Optional[Tuple[str, str]], + port: int, + host: str, + mmap: bool, + madv_random: bool, + parallel: int, + threads_each: int, + extra_args: List[str], + log_dir: Path, + ready_timeout: float, +) -> ServerProcess: + """Spawn `metagraph server_query`, return once it's accepting connections.""" + cmd: List[str] = [metagraph, "server_query"] + if graphs_csv is not None: + cmd.append(graphs_csv) + elif single_index is not None: + cmd += ["-i", single_index[0], "-a", single_index[1]] + else: + raise ValueError("Either graphs_csv or single_index must be provided.") + + cmd += ["--port", str(port)] + if host: + cmd += ["--address", host] + cmd += ["-p", str(parallel), "--threads-each", str(threads_each)] + if mmap: + cmd.append("--mmap") + if madv_random: + cmd.append("--madv-random") + cmd.append("-v") # required for trace lines to be emitted + cmd += extra_args + + log_path = log_dir / "server.log" + print(f"[bench] spawning: {' '.join(shlex.quote(c) for c in cmd)}", file=sys.stderr) + print(f"[bench] server log -> {log_path}", file=sys.stderr) + + log_file = open(log_path, "w", buffering=1) # line-buffered + # Note: server logs to stderr; merge stdout+stderr for completeness. + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=1, + text=True, + # New process group so SIGTERM only hits the server, not us. + preexec_fn=os.setsid if hasattr(os, "setsid") else None, + ) + + server = ServerProcess(proc=proc, log_path=log_path, drain_thread=None) # type: ignore[arg-type] + + def drain() -> None: + # Stream the server's combined stderr+stdout into both the in-memory + # buffer (for ready-detection and parsing) and the on-disk log. + assert proc.stdout is not None + for line in proc.stdout: + with server.log_lock: + server.log_buffer.append(line) + log_file.write(line) + log_file.close() + + drain_thread = threading.Thread(target=drain, daemon=True, name="server-drain") + drain_thread.start() + server.drain_thread = drain_thread + + # Multi-graph servers print "Ready to serve queries"; single-graph ones load + # asynchronously and only emit "Will listen on ...". For the latter we also + # need to retry the actual HTTP request because /search returns 503-ish + # while the lazy load is in flight; we handle that in send_query(). + ready_re = SERVER_READY_RE if graphs_csv is not None else SINGLE_GRAPH_READY_RE + deadline = time.monotonic() + ready_timeout + while time.monotonic() < deadline: + if proc.poll() is not None: + tail = server.snapshot()[-2000:] + raise RuntimeError( + f"Server exited (code {proc.returncode}) before ready.\n" + f"--- last log ---\n{tail}" + ) + if ready_re.search(server.snapshot()): + print(f"[bench] server ready on {host or '0.0.0.0'}:{port}", file=sys.stderr) + return server + time.sleep(0.2) + + server.stop() + raise TimeoutError( + f"Server did not become ready within {ready_timeout:.1f}s. " + f"See {log_path} for details." + ) + + +# --------------------------------------------------------------------------- +# Cache drop (best-effort; per-platform) +# --------------------------------------------------------------------------- + + +def drop_page_cache() -> Optional[str]: + """Try to drop the OS page cache. Returns None on success, else an error + message. Best-effort: silently no-ops if the script lacks privileges.""" + sys_name = platform.system() + try: + if sys_name == "Linux": + subprocess.run(["sync"], check=True) + with open("/proc/sys/vm/drop_caches", "w") as f: + f.write("3\n") + return None + elif sys_name == "Darwin": + r = subprocess.run(["purge"], check=False, capture_output=True) + if r.returncode != 0: + return f"`purge` failed: {r.stderr.decode().strip()}" + return None + else: + return f"unsupported OS: {sys_name}" + except PermissionError: + return "permission denied (run as root)" + except FileNotFoundError as e: + return f"missing tool: {e}" + except Exception as e: # pragma: no cover + return repr(e) + + +# --------------------------------------------------------------------------- +# Query +# --------------------------------------------------------------------------- + + +def fasta_payload(query_path: Path, max_seqs: Optional[int]) -> str: + """Read `query_path` (FASTA/FASTQ) and return its contents as FASTA text. + Optionally cap to first `max_seqs` records.""" + text = query_path.read_text() + if not max_seqs: + return text + out: List[str] = [] + seen = 0 + in_seq = False + for line in text.splitlines(): + if line.startswith(">"): + seen += 1 + if seen > max_seqs: + break + in_seq = True + out.append(line) + elif in_seq: + out.append(line) + return "\n".join(out) + "\n" + + +def send_query( + host: str, + port: int, + fasta: str, + discovery_fraction: float, + top_labels: int, + graph_names: Optional[List[str]], + timeout: float, + initial_retries: int, +) -> Tuple[float, int]: + """POST /search with `fasta`. Returns (wall_seconds, http_status).""" + payload: Dict[str, Any] = { + "FASTA": fasta, + "discovery_fraction": discovery_fraction, + "top_labels": top_labels, + "count_labels": True, + } + if graph_names is not None: + payload["graphs"] = graph_names + + body = json.dumps(payload).encode("utf-8") + req = urllib.request.Request( + f"http://{host}:{port}/search", + data=body, + method="POST", + headers={"Content-Type": "application/json"}, + ) + + last_err: Optional[Exception] = None + for attempt in range(initial_retries + 1): + t0 = time.monotonic() + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + resp.read() + return time.monotonic() - t0, resp.status + except urllib.error.HTTPError as e: + # 503 etc. while server is still warming up -- retry briefly. + last_err = e + if attempt < initial_retries: + time.sleep(0.5) + continue + raise + except (urllib.error.URLError, ConnectionResetError) as e: + last_err = e + if attempt < initial_retries: + time.sleep(0.5) + continue + raise + raise RuntimeError(f"unreachable: {last_err!r}") + + +# --------------------------------------------------------------------------- +# Log parsing +# --------------------------------------------------------------------------- + + +@dataclass +class KmerMapEvent: + threads: int + contigs: int + chunk_size: int + found: int + total: int + sec: float + + def to_dict(self) -> Dict[str, Any]: + return self.__dict__.copy() + + +@dataclass +class RdQueryEvent: + threads: int + rows_in: int + rows_out: int + expansion: float + traversal: float + call_rd_rows: float + set_bits: int + capacity: int + decoding: float + reconstruction: float + + def to_dict(self) -> Dict[str, Any]: + return self.__dict__.copy() + + +def parse_log(text: str) -> Tuple[List[KmerMapEvent], List[RdQueryEvent]]: + kmer_events: List[KmerMapEvent] = [] + rd_events: List[RdQueryEvent] = [] + for m in KMER_MAP_RE.finditer(text): + d = m.groupdict() + kmer_events.append(KmerMapEvent( + threads=int(d["threads"]), + contigs=int(d["contigs"]), + chunk_size=int(d["chunk_size"]), + found=int(d["found"]), + total=int(d["total"]), + sec=float(d["sec"]), + )) + for m in RD_QUERY_RE.finditer(text): + d = m.groupdict() + rd_events.append(RdQueryEvent( + threads=int(d["threads"]), + rows_in=int(d["rows_in"]), + rows_out=int(d["rows_out"]), + expansion=float(d["expansion"]), + traversal=float(d["traversal"]), + call_rd_rows=float(d["call_rd_rows"]), + set_bits=int(d["set_bits"]), + capacity=int(d["capacity"]), + decoding=float(d["decoding"]), + reconstruction=float(d["reconstruction"]), + )) + return kmer_events, rd_events + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + src = p.add_argument_group("server input (one required)") + src.add_argument("--graphs", type=Path, + help="CSV with rows ',,' for " + "multi-graph mode") + src.add_argument("--single-index", nargs=2, metavar=("GRAPH", "ANNO"), + help="single (graph, annotation) pair instead of a CSV") + + srv = p.add_argument_group("server") + srv.add_argument("--metagraph", required=True, + help="path to the metagraph binary") + srv.add_argument("--port", type=int, default=5555) + srv.add_argument("--host", default="127.0.0.1", + help="server bind address (default: 127.0.0.1)") + srv.add_argument("--mmap", action="store_true", + help="pass --mmap to server_query") + srv.add_argument("--madv-random", action="store_true", + help="pass --madv-random to server_query") + srv.add_argument("-p", "--parallel", type=int, default=1, + help="--parallel: max parallel connections") + srv.add_argument("--threads-each", type=int, default=1, + help="--threads-each: threads per graph") + srv.add_argument("--server-arg", action="append", default=[], + help="extra arg forwarded to server_query (repeatable)") + srv.add_argument("--ready-timeout", type=float, default=600.0, + help="seconds to wait for the server to become ready") + + qry = p.add_argument_group("query") + qry.add_argument("--query", required=True, type=Path, + help="FASTA file to send as the query body") + qry.add_argument("--graph-name", action="append", + help="restrict query to these graph names " + "(multi-graph mode; repeatable). Default: all.") + qry.add_argument("--discovery-fraction", type=float, default=0.0) + qry.add_argument("--top-labels", type=int, default=10000) + qry.add_argument("--max-seqs", type=int, default=0, + help="cap query to first N sequences (0 = all)") + qry.add_argument("--query-timeout", type=float, default=3600.0, + help="HTTP timeout per request") + + bnch = p.add_argument_group("benchmark") + bnch.add_argument("--warmup", type=int, default=0, + help="number of warmup queries (not counted)") + bnch.add_argument("--repeats", type=int, default=1, + help="number of timed queries") + bnch.add_argument("--drop-cache", action="store_true", + help="best-effort drop of OS page cache before each run " + "(needs root; warns if it fails)") + + out = p.add_argument_group("output") + out.add_argument("--json", type=Path, + help="write summary JSON to this path") + out.add_argument("--log-dir", type=Path, + help="dir to keep server logs (default: temp, deleted)") + + args = p.parse_args() + if (args.graphs is None) == (args.single_index is None): + p.error("exactly one of --graphs or --single-index is required") + return args + + +def main() -> int: + args = parse_args() + + if args.drop_cache: + err = drop_page_cache() + if err: + print(f"[bench] WARN: cache drop unavailable: {err}", file=sys.stderr) + + # --log-dir = persistent; otherwise a TemporaryDirectory we clean up. + if args.log_dir is not None: + args.log_dir.mkdir(parents=True, exist_ok=True) + log_dir_ctx: Any = _NoopCtx(args.log_dir) + else: + log_dir_ctx = tempfile.TemporaryDirectory(prefix="bench_metagraph_") + + with log_dir_ctx as log_dir_str: + log_dir = Path(log_dir_str) + + server = spawn_server( + metagraph=args.metagraph, + graphs_csv=str(args.graphs) if args.graphs else None, + single_index=tuple(args.single_index) if args.single_index else None, + port=args.port, + host=args.host, + mmap=args.mmap, + madv_random=args.madv_random, + parallel=args.parallel, + threads_each=args.threads_each, + extra_args=args.server_arg, + log_dir=log_dir, + ready_timeout=args.ready_timeout, + ) + + try: + fasta = fasta_payload(args.query, args.max_seqs) + print(f"[bench] query payload: {len(fasta)} bytes", file=sys.stderr) + + wall_times: List[float] = [] + log_marks: List[int] = [] # log offset before each TIMED request + + # Warmup + for i in range(args.warmup): + wall, status = send_query( + host=args.host, port=args.port, fasta=fasta, + discovery_fraction=args.discovery_fraction, + top_labels=args.top_labels, + graph_names=args.graph_name, + timeout=args.query_timeout, + initial_retries=20, # retry while async load completes + ) + print(f"[bench] warmup {i+1}/{args.warmup}: " + f"{wall:.2f}s (HTTP {status})", file=sys.stderr) + + # Timed runs + for i in range(args.repeats): + if args.drop_cache and i > 0: # cache already (just) dropped pre-spawn + err = drop_page_cache() + if err: + print(f"[bench] WARN: cache drop failed: {err}", + file=sys.stderr) + # Mark log offset so we can attribute trace lines to this run. + log_marks.append(len(server.snapshot())) + wall, status = send_query( + host=args.host, port=args.port, fasta=fasta, + discovery_fraction=args.discovery_fraction, + top_labels=args.top_labels, + graph_names=args.graph_name, + timeout=args.query_timeout, + initial_retries=2, + ) + wall_times.append(wall) + print(f"[bench] run {i+1}/{args.repeats}: " + f"{wall:.2f}s (HTTP {status})", file=sys.stderr) + + log_marks.append(len(server.snapshot())) + + finally: + # Give the server a moment to flush trailing trace lines, then stop. + time.sleep(0.5) + server.stop() + + # --- parse + report -------------------------------------------------- + full_log = server.snapshot() + + # Per-run aggregates: only count events that landed between consecutive marks. + per_run_kmer: List[List[KmerMapEvent]] = [] + per_run_rd: List[List[RdQueryEvent]] = [] + for i in range(args.repeats): + chunk = full_log[log_marks[i]:log_marks[i + 1]] + k_ev, r_ev = parse_log(chunk) + per_run_kmer.append(k_ev) + per_run_rd.append(r_ev) + + # Pooled across all timed runs. + all_kmer = [e for run in per_run_kmer for e in run] + all_rd = [e for run in per_run_rd for e in run] + + result = build_result(args, wall_times, all_kmer, all_rd, per_run_kmer, per_run_rd) + print_human_report(result) + + if args.json is not None: + args.json.write_text(json.dumps(result, indent=2, default=str) + "\n") + print(f"[bench] wrote {args.json}", file=sys.stderr) + + return 0 + + +def build_result( + args: argparse.Namespace, + wall_times: List[float], + all_kmer: List[KmerMapEvent], + all_rd: List[RdQueryEvent], + per_run_kmer: List[List[KmerMapEvent]], + per_run_rd: List[List[RdQueryEvent]], +) -> Dict[str, Any]: + return { + "config": { + "graphs": str(args.graphs) if args.graphs else None, + "single_index": list(args.single_index) if args.single_index else None, + "query": str(args.query), + "graph_names": args.graph_name, + "mmap": args.mmap, + "madv_random": args.madv_random, + "parallel": args.parallel, + "threads_each": args.threads_each, + "warmup": args.warmup, + "repeats": args.repeats, + "drop_cache": args.drop_cache, + "discovery_fraction": args.discovery_fraction, + "top_labels": args.top_labels, + }, + "client_wall_sec_per_run": wall_times, + "client_wall_sec": summarize(wall_times), + "kmer_mapping": { + "events_per_run": [len(r) for r in per_run_kmer], + "sec": summarize([e.sec for e in all_kmer]), + "contigs": summarize([e.contigs for e in all_kmer]), + "found_ratio": summarize([ + e.found / e.total if e.total else 0.0 for e in all_kmer + ]), + }, + "rd_query": { + "events_per_run": [len(r) for r in per_run_rd], + "rows_in": summarize([e.rows_in for e in all_rd]), + "rows_out": summarize([e.rows_out for e in all_rd]), + "expansion": summarize([e.expansion for e in all_rd]), + "traversal_sec": summarize([e.traversal for e in all_rd]), + "call_rd_rows_sec": summarize([e.call_rd_rows for e in all_rd]), + "decoding_sec": summarize([e.decoding for e in all_rd]), + "reconstruction_sec": summarize([e.reconstruction for e in all_rd]), + "set_bits": summarize([e.set_bits for e in all_rd]), + "capacity": summarize([e.capacity for e in all_rd]), + }, + "raw": { + "kmer_mapping_events": [e.to_dict() for e in all_kmer], + "rd_query_events": [e.to_dict() for e in all_rd], + }, + } + + +def print_human_report(r: Dict[str, Any]) -> None: + cfg = r["config"] + print() + print("=" * 72) + print("metagraph server_query benchmark") + print("=" * 72) + print(f" query: {cfg['query']}") + if cfg["graphs"]: + print(f" graphs csv: {cfg['graphs']}") + elif cfg["single_index"]: + print(f" index: {cfg['single_index'][0]}") + print(f" annotation: {cfg['single_index'][1]}") + if cfg["graph_names"]: + print(f" graph names: {', '.join(cfg['graph_names'])}") + print(f" flags: --parallel={cfg['parallel']} " + f"--threads-each={cfg['threads_each']} " + f"--mmap={cfg['mmap']} --madv-random={cfg['madv_random']}") + print(f" runs: warmup={cfg['warmup']} repeats={cfg['repeats']} " + f"drop_cache={cfg['drop_cache']}") + print() + + print("== Client wall time per request ==") + print(fmt_summary("wall", "s", r["client_wall_sec"])) + print() + + print("== K-mer mapping (Query graph construction / Contigs mapped) ==") + print(fmt_summary("time per event", "s", r["kmer_mapping"]["sec"])) + print(fmt_summary("contigs per event", "", r["kmer_mapping"]["contigs"])) + print(fmt_summary("found-ratio", "", r["kmer_mapping"]["found_ratio"])) + print(f" events per run: {r['kmer_mapping']['events_per_run']}") + print() + + print("== Annotation queries (RD query) ==") + print(fmt_summary("traversal time", "s", r["rd_query"]["traversal_sec"])) + print(fmt_summary("call_rd_rows time", "s", r["rd_query"]["call_rd_rows_sec"])) + print(fmt_summary("decoding time", "s", r["rd_query"]["decoding_sec"])) + print(fmt_summary("reconstruction time", "s", r["rd_query"]["reconstruction_sec"])) + print(fmt_summary("rows in", "", r["rd_query"]["rows_in"])) + print(fmt_summary("rows out", "", r["rd_query"]["rows_out"])) + print(fmt_summary("expansion", "x", r["rd_query"]["expansion"])) + print(f" events per run: {r['rd_query']['events_per_run']}") + print() + + +class _NoopCtx: + """Mimics tempfile.TemporaryDirectory's context but for a fixed path.""" + def __init__(self, p: Path) -> None: + self._p = str(p) + def __enter__(self) -> str: + return self._p + def __exit__(self, *exc: Any) -> None: + return None + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp index 913422decc..3488e91a43 100644 --- a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp +++ b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp @@ -2,11 +2,13 @@ #include #include +#include #include #include "common/algorithms.hpp" #include "common/serialization.hpp" +#include "common/utils/file_utils.hpp" #include "common/utils/template_utils.hpp" @@ -202,6 +204,8 @@ using NonzeroAndChildRows = std::pair, std::vector& row_ids, ThreadPool* thread_pool, bool adaptive_chunk_size) const { + prefetch_if_dense(row_ids.size()); + std::vector nonzero_indices; std::vector child_row_ids; nonzero_indices.reserve(row_ids.size()); @@ -468,16 +472,44 @@ std::vector BRWT::get_column(Column column) const { return rows; } +void BRWT::prefetch_if_dense(size_t num_queries) const { + if (!utils::with_madvise() || !nonzero_rows_mmap_addr_ || !nonzero_rows_mmap_size_) + return; + const long psz_l = sysconf(_SC_PAGESIZE); + if (psz_l <= 0 || num_queries == 0) + return; + const size_t pagesize = static_cast(psz_l); + if (num_queries > SIZE_MAX / pagesize) + return; + const size_t q_span = num_queries * pagesize; + constexpr unsigned ALPHA_NUM = 1; + constexpr unsigned ALPHA_DEN = 10; + if (q_span > SIZE_MAX / ALPHA_DEN) + return; + if (q_span * ALPHA_DEN < nonzero_rows_mmap_size_ * ALPHA_NUM) + return; + utils::madvise_willneed(nonzero_rows_mmap_addr_, nonzero_rows_mmap_size_); +} + bool BRWT::load(std::istream &in) { if (!in.good()) return false; try { + nonzero_rows_mmap_addr_ = nullptr; + nonzero_rows_mmap_size_ = 0; + if (!assignments_.load(in)) return false; + const auto nonzero_start = static_cast(in.tellg()); if (!nonzero_rows_->load(in)) return false; + if (void *base = utils::get_mmap_data(in, nonzero_start)) { + const auto nonzero_end = static_cast(in.tellg()); + nonzero_rows_mmap_addr_ = base; + nonzero_rows_mmap_size_ = static_cast(nonzero_end - nonzero_start); + } size_t num_child_nodes = load_number(in); child_nodes_.clear(); diff --git a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp index 6d83fd1724..cf59d59e95 100644 --- a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp +++ b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp @@ -44,6 +44,10 @@ class BRWT : public BinaryMatrix, public GetEntrySupport { bool load(std::istream &in) override; void serialize(std::ostream &out) const override; + // Hint MADV_WILLNEED on this node's nonzero_rows_ mmap when |num_queries| + // is large enough to amortize (threshold ~10% of bitmap span in pages). + void prefetch_if_dense(size_t num_queries) const; + // number of ones in the matrix uint64_t num_relations() const override; @@ -88,6 +92,9 @@ class BRWT : public BinaryMatrix, public GetEntrySupport { std::unique_ptr nonzero_rows_; // generally, these child matrices can be abstract BinaryMatrix instances std::vector> child_nodes_; + + void *nonzero_rows_mmap_addr_ = nullptr; + size_t nonzero_rows_mmap_size_ = 0; }; } // namespace matrix diff --git a/metagraph/src/annotation/binary_matrix/row_diff/row_diff.cpp b/metagraph/src/annotation/binary_matrix/row_diff/row_diff.cpp index 7668b11bcb..cb2e780850 100644 --- a/metagraph/src/annotation/binary_matrix/row_diff/row_diff.cpp +++ b/metagraph/src/annotation/binary_matrix/row_diff/row_diff.cpp @@ -42,31 +42,52 @@ node_index row_diff_successor(const graph::DeBruijnGraph &graph, namespace matrix { void IRowDiff::load_anchor(const std::string &filename) { + anchor_mmap_addr_ = nullptr; + anchor_mmap_size_ = 0; std::unique_ptr f = utils::open_ifstream(filename); if (!f->good()) { logger->error("Cannot open anchor file '{}': {}", filename, utils::file_read_failure_detail(filename)); std::exit(1); } + const auto anchor_start = static_cast(f->tellg()); if (!anchor_.load(*f)) { logger->error("Cannot load anchor from '{}': {}", filename, utils::file_read_failure_detail(filename)); std::exit(1); } + if (void *base = utils::get_mmap_data(*f, anchor_start)) { + const auto anchor_end = static_cast(f->tellg()); + anchor_mmap_addr_ = base; + anchor_mmap_size_ = static_cast(anchor_end - anchor_start); + } } void IRowDiff::load_fork_succ(const std::string &filename) { + fork_succ_mmap_addr_ = nullptr; + fork_succ_mmap_size_ = 0; std::unique_ptr f = utils::open_ifstream(filename); if (!f->good()) { logger->error("Cannot open fork successor file '{}': {}", filename, utils::file_read_failure_detail(filename)); std::exit(1); } + const auto fork_succ_start = static_cast(f->tellg()); if (!fork_succ_.load(*f)) { logger->error("Cannot load fork successor bitmap from '{}': {}", filename, utils::file_read_failure_detail(filename)); std::exit(1); } + if (void *base = utils::get_mmap_data(*f, fork_succ_start)) { + const auto fork_succ_end = static_cast(f->tellg()); + fork_succ_mmap_addr_ = base; + fork_succ_mmap_size_ = static_cast(fork_succ_end - fork_succ_start); + } +} + +void IRowDiff::prefetch() const { + utils::madvise_willneed(anchor_mmap_addr_, anchor_mmap_size_); + utils::madvise_willneed(fork_succ_mmap_addr_, fork_succ_mmap_size_); } std::tuple, diff --git a/metagraph/src/annotation/binary_matrix/row_diff/row_diff.hpp b/metagraph/src/annotation/binary_matrix/row_diff/row_diff.hpp index d9a9808314..7068990ffb 100644 --- a/metagraph/src/annotation/binary_matrix/row_diff/row_diff.hpp +++ b/metagraph/src/annotation/binary_matrix/row_diff/row_diff.hpp @@ -51,6 +51,9 @@ class IRowDiff { const fork_succ_bv_type& fork_succ() const { return fork_succ_; } + // Hint MADV_WILLNEED on mmap-backed anchor + fork_succ bitmaps (no-op if not mmap). + void prefetch() const; + protected: // get row-diff paths starting at |row_ids| // Returns: (rd_ids, rd_paths_trunc, times_traversed, groups) @@ -70,6 +73,11 @@ class IRowDiff { const graph::DeBruijnGraph *graph_ = nullptr; anchor_bv_type anchor_; fork_succ_bv_type fork_succ_; + + void *anchor_mmap_addr_ = nullptr; + size_t anchor_mmap_size_ = 0; + void *fork_succ_mmap_addr_ = nullptr; + size_t fork_succ_mmap_size_ = 0; }; /** diff --git a/metagraph/src/annotation/binary_matrix/row_disk/row_disk.cpp b/metagraph/src/annotation/binary_matrix/row_disk/row_disk.cpp index b182b9d726..2a63d33f8f 100644 --- a/metagraph/src/annotation/binary_matrix/row_disk/row_disk.cpp +++ b/metagraph/src/annotation/binary_matrix/row_disk/row_disk.cpp @@ -11,8 +11,13 @@ namespace matrix { using mtg::common::logger; +void RowDisk::prefetch_boundary() const { + utils::madvise_willneed(boundary_mmap_addr_, boundary_mmap_size_); +} + std::vector RowDisk::get_rows(const std::vector &row_ids) const { + prefetch_boundary(); View view = get_view(); std::vector rows(row_ids.size()); for (size_t i = 0; i < row_ids.size(); ++i) { @@ -61,9 +66,20 @@ bool RowDisk::load(std::istream &f) { assert(boundary_start >= buffer_params_.offset); iv_size_on_disk_ = boundary_start - buffer_params_.offset; + boundary_mmap_addr_ = nullptr; + boundary_mmap_size_ = 0; // boundary_ is too large to load into RAM, always mmap it. utils::load_mmap_random(buffer_params_.filename, boundary_start, - [&](std::istream &in) { boundary_.load(in); }); + [&](std::istream &in) { + const auto boundary_byte_start = static_cast(in.tellg()); + boundary_.load(in); + if (void *base = utils::get_mmap_data(in, boundary_byte_start)) { + const auto boundary_byte_end = static_cast(in.tellg()); + boundary_mmap_addr_ = base; + boundary_mmap_size_ + = static_cast(boundary_byte_end - boundary_byte_start); + } + }); num_rows_ = boundary_.num_set_bits(); num_relations_ = boundary_.size() - num_rows_; diff --git a/metagraph/src/annotation/binary_matrix/row_disk/row_disk.hpp b/metagraph/src/annotation/binary_matrix/row_disk/row_disk.hpp index 9b3429c441..c67a09cb5d 100644 --- a/metagraph/src/annotation/binary_matrix/row_disk/row_disk.hpp +++ b/metagraph/src/annotation/binary_matrix/row_disk/row_disk.hpp @@ -44,6 +44,7 @@ class RowDisk : public BinaryMatrix { const bit_vector_small& get_boundary() const { return boundary_; } private: + void prefetch_boundary() const; // For the multithreading to work properly, we open int_vector_buffer<> in // a special View class that has an actual implementation of the method. class View { @@ -80,6 +81,9 @@ class RowDisk : public BinaryMatrix { uint64_t num_relations_ = 0; size_t iv_size_on_disk_ = 0; // for non-static serialization + + void *boundary_mmap_addr_ = nullptr; + size_t boundary_mmap_size_ = 0; }; } // namespace matrix diff --git a/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.cpp b/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.cpp index e35a05604d..ebeda3de57 100644 --- a/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.cpp +++ b/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.cpp @@ -9,8 +9,13 @@ namespace matrix { using mtg::common::logger; +void CoordRowDisk::prefetch_boundary() const { + utils::madvise_willneed(boundary_mmap_addr_, boundary_mmap_size_); +} + std::vector CoordRowDisk::get_rows(const std::vector &row_ids) const { + prefetch_boundary(); View view = get_view(); std::vector rows(row_ids.size()); for (size_t i = 0; i < row_ids.size(); ++i) { @@ -21,12 +26,14 @@ CoordRowDisk::get_rows(const std::vector &row_ids) const { std::vector CoordRowDisk::get_row_values(const std::vector &rows, size_t num_threads) const { + prefetch_boundary(); return get_row_data_parallel(rows, num_threads, [&](const auto &rows) { return get_view().get_row_values(rows); }); } std::vector CoordRowDisk::get_row_tuples(const std::vector &rows, size_t num_threads) const { + prefetch_boundary(); return get_row_data_parallel(rows, num_threads, [&](const auto &rows) { return get_view().get_row_tuples(rows); }); } @@ -161,10 +168,19 @@ bool CoordRowDisk::load(std::istream &f) { assert(boundary_start >= buffer_params_.offset); + boundary_mmap_addr_ = nullptr; + boundary_mmap_size_ = 0; // boundary_ is too large to load into RAM, always mmap it. utils::load_mmap_random(buffer_params_.filename, boundary_start, [&](std::istream &in) { + const auto boundary_byte_start = static_cast(in.tellg()); boundary_.load(in); + if (void *base = utils::get_mmap_data(in, boundary_byte_start)) { + const auto boundary_byte_end = static_cast(in.tellg()); + boundary_mmap_addr_ = base; + boundary_mmap_size_ + = static_cast(boundary_byte_end - boundary_byte_start); + } num_attributes_ = load_number(in); }); diff --git a/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.hpp b/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.hpp index 3f1a5a8b10..3fe10ed0fe 100644 --- a/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.hpp +++ b/metagraph/src/annotation/int_matrix/row_disk/coord_row_disk.hpp @@ -63,6 +63,7 @@ class CoordRowDisk : public BinaryMatrix, public MultiIntMatrix { const BinaryMatrix& get_binary_matrix() const { return *this; } private: + void prefetch_boundary() const; // For the multithreading to work properly, we open int_vector_buffer<> in // a special View class that has an actual implementation of the method. class View { @@ -129,6 +130,9 @@ class CoordRowDisk : public BinaryMatrix, public MultiIntMatrix { uint64_t bits_for_number_of_vals_ = 0; uint64_t bits_for_single_value_ = 0; uint64_t num_rows_ = 0; + + void *boundary_mmap_addr_ = nullptr; + size_t boundary_mmap_size_ = 0; }; } // namespace matrix diff --git a/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.cpp b/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.cpp index 2a88493917..be03e2ce59 100644 --- a/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.cpp +++ b/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.cpp @@ -9,8 +9,13 @@ namespace matrix { using mtg::common::logger; +void IntRowDisk::prefetch_boundary() const { + utils::madvise_willneed(boundary_mmap_addr_, boundary_mmap_size_); +} + std::vector IntRowDisk::get_rows(const std::vector &row_ids) const { + prefetch_boundary(); View view = get_view(); std::vector rows(row_ids.size()); for (size_t i = 0; i < row_ids.size(); ++i) { @@ -21,6 +26,7 @@ IntRowDisk::get_rows(const std::vector &row_ids) const { std::vector IntRowDisk::get_row_values(const std::vector &rows, size_t num_threads) const { + prefetch_boundary(); return get_row_data_parallel(rows, num_threads, [&](const auto &rows) { return get_view().get_row_values(rows); }); } @@ -107,9 +113,20 @@ bool IntRowDisk::load(std::istream &f) { assert(boundary_start >= buffer_params_.offset); + boundary_mmap_addr_ = nullptr; + boundary_mmap_size_ = 0; // boundary_ is too large to load into RAM, always mmap it. utils::load_mmap_random(buffer_params_.filename, boundary_start, - [&](std::istream &in) { boundary_.load(in); }); + [&](std::istream &in) { + const auto boundary_byte_start = static_cast(in.tellg()); + boundary_.load(in); + if (void *base = utils::get_mmap_data(in, boundary_byte_start)) { + const auto boundary_byte_end = static_cast(in.tellg()); + boundary_mmap_addr_ = base; + boundary_mmap_size_ + = static_cast(boundary_byte_end - boundary_byte_start); + } + }); num_rows_ = boundary_.num_set_bits(); diff --git a/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.hpp b/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.hpp index 4148acb40a..8571c41055 100644 --- a/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.hpp +++ b/metagraph/src/annotation/int_matrix/row_disk/int_row_disk.hpp @@ -54,6 +54,7 @@ class IntRowDisk : public BinaryMatrix, public IntMatrix { const BinaryMatrix& get_binary_matrix() const { return *this; } private: + void prefetch_boundary() const; // For the multithreading to work properly, we open int_vector_buffer<> in // a special View class that has an actual implementation of the method. class View { @@ -111,6 +112,9 @@ class IntRowDisk : public BinaryMatrix, public IntMatrix { uint64_t bits_for_col_id_ = 0; uint64_t bits_for_value_ = 0; uint64_t num_rows_ = 0; + + void *boundary_mmap_addr_ = nullptr; + size_t boundary_mmap_size_ = 0; }; } // namespace matrix diff --git a/metagraph/src/cli/query.cpp b/metagraph/src/cli/query.cpp index db4da5f32d..d2d5a2d78e 100644 --- a/metagraph/src/cli/query.cpp +++ b/metagraph/src/cli/query.cpp @@ -12,8 +12,10 @@ #include "common/utils/template_utils.hpp" #include "common/threads/threading.hpp" #include "common/vectors/vector_algorithm.hpp" +#include "annotation/binary_matrix/row_diff/row_diff.hpp" #include "annotation/representation/annotation_matrix/static_annotators_def.hpp" #include "graph/alignment/dbg_aligner.hpp" +#include "graph/representation/canonical_dbg.hpp" #include "graph/representation/hash/dbg_hash_ordered.hpp" #include "graph/representation/succinct/dbg_succinct.hpp" #include "graph/representation/succinct/boss_construct.hpp" @@ -1223,6 +1225,24 @@ size_t query_fasta(const std::string &file, const graph::align::DBGAlignerConfig *aligner_config) { logger->trace("Parsing sequences from file '{}'", file); + // Warm the suffix-ranges index pages before processing this file so that + // page faults (from mmap eviction under memory pressure) are amortized + // by an async prefetch instead of stalling individual k-mer lookups. + // Handles the PRIMARY graph case where DBGSuccinct is wrapped in CanonicalDBG. + { + const DeBruijnGraph *graph = &anno_graph.get_graph(); + if (const auto *canonical = dynamic_cast(graph)) + graph = &canonical->get_graph(); + if (const auto *dbg_succ = dynamic_cast(graph)) { + dbg_succ->prefetch_suffix_ranges(); + dbg_succ->prefetch_bloom_filter(); + } + } + + if (const auto *rd = dynamic_cast( + &anno_graph.get_annotator().get_matrix())) + rd->prefetch(); + seq_io::FastaParser fasta_parser(file, config.forward_and_reverse); // Only query_coords/count_kmers if using coord/count aware index. diff --git a/metagraph/src/common/utils/file_utils.cpp b/metagraph/src/common/utils/file_utils.cpp index e80617bad0..1a0475e91b 100644 --- a/metagraph/src/common/utils/file_utils.cpp +++ b/metagraph/src/common/utils/file_utils.cpp @@ -108,6 +108,26 @@ void madvise_random_range(std::istream &f, std::streamoff start, std::streamoff } } +void madvise_willneed(void *addr, size_t length) { + if (!with_madvise() || !addr || !length) + return; + const size_t pagesize = sysconf(_SC_PAGESIZE); + auto raw = reinterpret_cast(addr); + auto aligned = raw & ~(pagesize - 1u); + size_t total = length + (raw - aligned); + if (madvise(reinterpret_cast(aligned), total, MADV_WILLNEED)) { + logger->warn("madvise(MADV_WILLNEED) failed for [{}, {})", + addr, static_cast(addr) + length); + } +} + +void *get_mmap_data(std::istream &f, std::streamoff offset) { + auto *mmap_in = dynamic_cast(&f); + if (!mmap_in) + return nullptr; + return mmap_in->get_mmap_context()->data() + offset; +} + void load_mmap_random(const std::string &filename, std::streamoff offset, const std::function &fn) { sdsl::mmap_ifstream in(filename); diff --git a/metagraph/src/common/utils/file_utils.hpp b/metagraph/src/common/utils/file_utils.hpp index 0035f8bcf5..f4b6cf727c 100644 --- a/metagraph/src/common/utils/file_utils.hpp +++ b/metagraph/src/common/utils/file_utils.hpp @@ -62,6 +62,18 @@ void madvise_random_range(std::istream &f, std::streamoff start = 0, std::streamoff length = -1); +// Hint MADV_WILLNEED on `[addr, addr + length)`, triggering async prefetch +// of those pages into the page cache. Page-aligns the start address. +// No-op when `with_madvise()` is false or `addr` is null. +void madvise_willneed(void *addr, size_t length); + +// Returns a pointer into the mmap region backing `f` at byte `offset` from +// the start of the file, or `nullptr` if `f` is not an `sdsl::mmap_ifstream`. +// The pointer remains valid as long as the underlying mmap context is alive +// (the context is shared with sdsl objects loaded from `f`, so it outlives +// `f` itself when those objects keep a reference to it). +void *get_mmap_data(std::istream &f, std::streamoff offset = 0); + // Open `filename` as an `sdsl::mmap_ifstream` seeked to `offset`, invoke // `fn` so the caller can load from it, then (if `with_madvise()`) hint // MADV_RANDOM on the whole file mapping. diff --git a/metagraph/src/graph/representation/succinct/dbg_succinct.cpp b/metagraph/src/graph/representation/succinct/dbg_succinct.cpp index d3ddf10ea9..892db2cd21 100644 --- a/metagraph/src/graph/representation/succinct/dbg_succinct.cpp +++ b/metagraph/src/graph/representation/succinct/dbg_succinct.cpp @@ -687,6 +687,15 @@ uint64_t DBGSuccinct::max_index() const { return boss_graph_->num_edges(); } +void DBGSuccinct::prefetch_suffix_ranges() const { + utils::madvise_willneed(suffix_ranges_mmap_addr_, suffix_ranges_mmap_size_); +} + +void DBGSuccinct::prefetch_bloom_filter() const { + if (bloom_filter_) + bloom_filter_->prefetch(); +} + bool DBGSuccinct::load_without_mask(const std::string &filename) { // release the old mask valid_edges_.reset(); @@ -700,8 +709,20 @@ bool DBGSuccinct::load_without_mask(const std::string &filename) { mode_ = static_cast(load_number(*in)); - if (!boss_graph_->load_suffix_ranges(*in)) + suffix_ranges_mmap_addr_ = nullptr; + suffix_ranges_mmap_size_ = 0; + auto suffix_ranges_start = static_cast(in->tellg()); + if (!boss_graph_->load_suffix_ranges(*in)) { logger->warn("No index for node ranges could be loaded"); + } else if (void *base = utils::get_mmap_data(*in, suffix_ranges_start)) { + // sdsl zero-copy: when loaded from an mmap stream, the suffix-ranges + // data lives in the mmap region. Record its address+size so that + // prefetch_suffix_ranges() can warm those pages later. + auto suffix_ranges_end = static_cast(in->tellg()); + suffix_ranges_mmap_addr_ = base; + suffix_ranges_mmap_size_ + = static_cast(suffix_ranges_end - suffix_ranges_start); + } // hint random access for query-time traversal of the loaded data utils::madvise_random_range(*in); diff --git a/metagraph/src/graph/representation/succinct/dbg_succinct.hpp b/metagraph/src/graph/representation/succinct/dbg_succinct.hpp index 2f243189f4..3f182fb98a 100644 --- a/metagraph/src/graph/representation/succinct/dbg_succinct.hpp +++ b/metagraph/src/graph/representation/succinct/dbg_succinct.hpp @@ -173,6 +173,13 @@ class DBGSuccinct : public DeBruijnGraph { virtual boss::BOSS& get_boss() final { return *boss_graph_; } virtual boss::BOSS* release_boss() final { return boss_graph_.release(); } + // Issue madvise(WILLNEED) on the suffix-ranges index pages so they are + // (asynchronously) fetched into the page cache before subsequent queries. + // No-op when the graph was not loaded with mmap or madvise is disabled. + void prefetch_suffix_ranges() const; + + void prefetch_bloom_filter() const; + virtual bool operator==(const DeBruijnGraph &other) const override final; virtual const std::string& alphabet() const override final; @@ -207,6 +214,12 @@ class DBGSuccinct : public DeBruijnGraph { std::unique_ptr> bloom_filter_; + // Virtual address range of the suffix_ranges index inside the mmap'd + // graph file. Set in load_without_mask() when loaded via mmap; remains + // valid as long as this object is alive (sdsl keeps the mapping open). + void *suffix_ranges_mmap_addr_ = nullptr; + size_t suffix_ranges_mmap_size_ = 0; + std::unique_ptr generate_valid_kmer_mask(size_t num_threads, bool with_pruning) const; }; diff --git a/metagraph/src/kmer/kmer_bloom_filter.cpp b/metagraph/src/kmer/kmer_bloom_filter.cpp index 2f277900b1..b4242bc232 100644 --- a/metagraph/src/kmer/kmer_bloom_filter.cpp +++ b/metagraph/src/kmer/kmer_bloom_filter.cpp @@ -195,20 +195,36 @@ ::serialize(std::ostream &out) const { template bool KmerBloomFilter ::load(std::istream &in) { + bloom_mmap_addr_ = nullptr; + bloom_mmap_size_ = 0; if (!in.good()) return false; try { + const auto bloom_start = static_cast(in.tellg()); k_ = load_number(in); canonical_mode_ = load_number(in); const_cast(hasher_) = KmerHasher(k_); - return filter_.load(in); + if (!filter_.load(in)) + return false; + + if (void *base = utils::get_mmap_data(in, bloom_start)) { + const auto bloom_end = static_cast(in.tellg()); + bloom_mmap_addr_ = base; + bloom_mmap_size_ = static_cast(bloom_end - bloom_start); + } + return true; } catch (...) { return false; } } +template +void KmerBloomFilter::prefetch() const { + utils::madvise_willneed(bloom_mmap_addr_, bloom_mmap_size_); +} + template class KmerBloomFilter<>; diff --git a/metagraph/src/kmer/kmer_bloom_filter.hpp b/metagraph/src/kmer/kmer_bloom_filter.hpp index 5bfc4ac0a8..d0c45138dc 100644 --- a/metagraph/src/kmer/kmer_bloom_filter.hpp +++ b/metagraph/src/kmer/kmer_bloom_filter.hpp @@ -3,6 +3,7 @@ #include "common/hashers/rolling_hasher.hpp" #include "common/bloom_filter.hpp" +#include "common/utils/file_utils.hpp" namespace mtg { @@ -53,11 +54,16 @@ class KmerBloomFilter { const BloomFilter& get_filter() const { return filter_; } + void prefetch() const; + private: BloomFilter filter_; bool canonical_mode_; size_t k_; const KmerHasher hasher_; + + void *bloom_mmap_addr_ = nullptr; + size_t bloom_mmap_size_ = 0; };