diff --git a/packages/python/tests/e2e/speed_of_light.py b/packages/python/tests/e2e/speed_of_light.py new file mode 100644 index 000000000..66173c7f1 --- /dev/null +++ b/packages/python/tests/e2e/speed_of_light.py @@ -0,0 +1,124 @@ +"""Harness for the use-case speed-of-light tests. + +Each use-case test asks one question of a generated corpus twice: once through +the fastest native pipeline (find, sort, join, awk, a single parser process) +and once through the real `dirsql` launcher. The rows must be identical and +dirsql's wall time, after subtracting the launcher's measured startup, must be +within ten percent of native. The bar is a ratio, never a millisecond number, +so it holds on any machine. + +Startup is the mean of ten `dirsql --help` runs; `--help` and `SELECT 1` cost +the same, so it stands in for the fixed per-invocation cost. A corpus too +small to time is grown by doubling until the native pipeline takes at least a +second, so the comparison is above the noise floor. Each side is run three +times and the fastest run counts. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import time + +ONE_SECOND = 1.0 +TOLERANCE = 1.1 +STARTUP_PROBES = 10 +BEST_OF = 3 + + +def cli() -> str: + dirsql = shutil.which("dirsql") + assert dirsql is not None, ( + "`dirsql` console script not on PATH -- run `uv run maturin develop`" + ) + return dirsql + + +def timed(argv, cwd, timeout, env=None): + started = time.perf_counter() + proc = subprocess.run( + argv, + cwd=str(cwd), + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=timeout, + env=env, + check=False, + ) + return proc, time.perf_counter() - started + + +def startup_seconds(argv=None, cwd=None) -> float: + argv = [cli(), "--help"] if argv is None else argv + samples = [] + for _ in range(STARTUP_PROBES): + proc, seconds = timed(argv, cwd or ".", timeout=60) + assert proc.returncode == 0, proc.stderr + samples.append(seconds) + return sum(samples) / len(samples) + + +def dirsql_rows(proc, columns): + assert proc.returncode == 0, ( + f"dirsql failed: stdout={proc.stdout!r} stderr={proc.stderr!r}" + ) + return [tuple(row[column] for column in columns) for row in json.loads(proc.stdout)] + + +def grow_until_native_takes_a_second(build, native, start, ceiling): + """Grow the corpus by doubling until `native` runs for a second. + + `build(lo, hi)` adds items `lo` through `hi - 1` to the corpus already + holding `lo` items, so each doubling writes only the new half. + """ + built, size = 0, start + while True: + build(built, size) + built = size + _, seconds = native() + if seconds >= ONE_SECOND or size >= ceiling: + return size + size *= 2 + + +def best_of(run, name, hopeless=None): + """Run `run` three times; return the rows and wall time of the fastest. + + `run` returns `(rows, seconds)`. A run that raises `TimeoutExpired` fails + the test immediately with the timeout it blew through. A run slower than + `hopeless` seconds ends the series early: noise does not account for an + order of magnitude, so no retry could bring the best under the bar. + """ + fastest = None + for _ in range(BEST_OF): + try: + rows, seconds = run() + except subprocess.TimeoutExpired as expired: + raise AssertionError( + f"{name} did not finish within {expired.timeout:.0f}s" + ) from None + if fastest is None or seconds < fastest[1]: + fastest = (rows, seconds) + if hopeless is not None and seconds > hopeless: + break + return fastest + + +def hopeless_seconds(native_seconds: float, startup: float) -> float: + return startup + 10 * TOLERANCE * native_seconds + + +def dirsql_timeout(native_seconds: float) -> float: + return max(60.0, 30 * native_seconds) + + +def assert_speed_of_light(name, native_seconds, dirsql_seconds, startup): + measured = dirsql_seconds - startup + budget = TOLERANCE * native_seconds + assert measured <= budget, ( + f"{name}: dirsql {dirsql_seconds:.3f}s minus startup {startup:.3f}s" + f" = {measured:.3f}s, over {TOLERANCE}x native {native_seconds:.3f}s" + f" = {budget:.3f}s" + ) diff --git a/packages/python/tests/e2e/use_case_dossier_test.py b/packages/python/tests/e2e/use_case_dossier_test.py new file mode 100644 index 000000000..c413e0409 --- /dev/null +++ b/packages/python/tests/e2e/use_case_dossier_test.py @@ -0,0 +1,153 @@ +"""Use case 1, the dossier: three selections over a note tree, UNIONed. + +A standing set of background files, the five newest dailies by name, and the +ten most recently modified markdown files outside the dailies, in one ranked +list. Native is three `find` runs with `-printf`, each sorted and cut, then +one final sort. dirsql is the article's query through the real launcher. +No mocks: real console script, real process, real filesystem. +""" + +from __future__ import annotations + +import os +import shutil +import time + +import pytest + +from .speed_of_light import ( + assert_speed_of_light, + best_of, + cli, + dirsql_rows, + dirsql_timeout, + grow_until_native_takes_a_second, + hopeless_seconds, + startup_seconds, + timed, +) + +DAY = 86400 +WINDOW_DAYS = 14 +STANDING = ("situation.md", "projects.md", "profile.md") +COLUMNS = ("rank", "kind", "path", "mtime") + +QUERY = """ +WITH standing AS ( + SELECT path, mtime FROM './' + WHERE dir = 'planner/background' + AND basename IN ('situation.md', 'projects.md', 'profile.md') +), +dailies AS ( + SELECT path, mtime FROM './' + WHERE path LIKE 'vault/Dailies/Daily %' + ORDER BY basename DESC LIMIT 5 +), +recent AS ( + SELECT path, mtime FROM './' + WHERE ext = 'md' + AND path NOT LIKE 'vault/Dailies/%' + AND mtime > strftime('%s', 'now') - 14 * 86400 + ORDER BY mtime DESC LIMIT 10 +) +SELECT 1 AS rank, 'standing' AS kind, path, mtime FROM standing +UNION ALL SELECT 2, 'daily', path, mtime FROM dailies +UNION ALL SELECT 3, 'recent', path, mtime FROM recent +ORDER BY rank, mtime DESC, path +""" + +NATIVE = r""" +set -e +export LC_ALL=C +cutoff=$(( $(date +%s) - 14 * 86400 )) +{ + find planner/background -maxdepth 1 \ + \( -name situation.md -o -name projects.md -o -name profile.md \) \ + -printf '1\tstanding\t%p\t%T@\n' + find vault/Dailies -name 'Daily *' -printf '%f\t%p\t%T@\n' \ + | sort -t "$(printf '\t')" -k1,1r | head -5 \ + | awk -F '\t' -v OFS='\t' '{ print 2, "daily", $2, $3 }' + find . -path ./vault/Dailies -prune -o -name '*.md' -newermt "@$cutoff" \ + -printf '%T@\t%P\n' \ + | sort -rn | head -10 \ + | awk -F '\t' -v OFS='\t' '{ print 3, "recent", $2, $1 }' +} | sort -t "$(printf '\t')" -k1,1n -k4,4rn -k3,3 +""" + + +def _write(path, mtime, body="# note\nbody\n"): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(body) + os.utime(path, (mtime, mtime)) + + +def build_tree(root, lo, hi): + """Bulk notes `lo` through `hi - 1`: the first thousand inside the recency + window, the rest older, every tenth one a `.txt` the ext filter must + drop. Mtimes are distinct integers nowhere near the window's edge.""" + now = int(time.time()) + if lo == 0: + for name in STANDING: + _write(root / "planner" / "background" / name, now - 70 * DAY) + for day in range(1, 41): + stamp = now - day * DAY + date = time.strftime("%Y.%m.%d", time.gmtime(stamp)) + _write(root / "vault" / "Dailies" / f"Daily {date}.md", stamp) + for i in range(lo, hi): + ext = "txt" if i % 10 == 0 else "md" + mtime = now - 3600 - i * 600 if i < 1000 else now - 15 * DAY - i * 7 + _write( + root / "planner" / "internal" / "notes" / str(i % 1000) / f"note-{i}.{ext}", + mtime, + ) + + +def native_rows(proc): + assert proc.returncode == 0, proc.stderr + rows = [] + for line in proc.stdout.splitlines(): + rank, kind, path, mtime = line.split("\t") + rows.append((int(rank), kind, path, int(float(mtime)))) + return rows + + +def describe_dossier_speed_of_light(): + @pytest.fixture + def root(tmp_path): + tree = tmp_path / "notes" + try: + yield tree + finally: + shutil.rmtree(tree, ignore_errors=True) + + @pytest.mark.xfail( + strict=True, + reason="each CTE's path-table reference walks the tree again, and one walk is slower than find", + ) + def it_matches_native_rows_within_the_bar(root): + startup = startup_seconds() + + def native(): + proc, seconds = timed(["bash", "-c", NATIVE], root, timeout=600) + return native_rows(proc), seconds + + grow_until_native_takes_a_second( + lambda lo, hi: build_tree(root, lo, hi), native, start=50_000, ceiling=2**21 + ) + expected, native_seconds = best_of(native, "native") + assert len(expected) == 18, expected + + def dirsql(): + proc, seconds = timed( + [cli(), "query", QUERY], + root, + timeout=dirsql_timeout(native_seconds), + ) + return dirsql_rows(proc, COLUMNS), seconds + + actual, dirsql_seconds = best_of( + dirsql, "dirsql", hopeless_seconds(native_seconds, startup) + ) + + assert actual == expected + assert_speed_of_light("dossier", native_seconds, dirsql_seconds, startup) diff --git a/packages/python/tests/e2e/use_case_join_test.py b/packages/python/tests/e2e/use_case_join_test.py new file mode 100644 index 000000000..415269955 --- /dev/null +++ b/packages/python/tests/e2e/use_case_join_test.py @@ -0,0 +1,117 @@ +"""Use case 2, the arXiv sidecar layout: a join of two path tables. + +One folder per paper holding `abstract.md` and `title.md`; the question is +every paper that has both, with its title and the abstract's size. Native is +two `find` listings, one through `awk` for the title text, sorted and joined +with `join(1)`. dirsql joins the two path tables on `dir` through the real +launcher. No mocks: real console script, real process, real filesystem. +""" + +from __future__ import annotations + +import shutil + +import pytest + +from .speed_of_light import ( + assert_speed_of_light, + best_of, + cli, + dirsql_rows, + dirsql_timeout, + grow_until_native_takes_a_second, + hopeless_seconds, + startup_seconds, + timed, +) + +COLUMNS = ("dir", "title", "abstract_bytes") +QUERY = """ +SELECT a.dir, t.content AS title, a.size AS abstract_bytes +FROM './papers/*/abstract.md' a +JOIN './papers/*/title.md' t ON t.dir = a.dir +ORDER BY a.dir +""" + +NATIVE = r""" +set -e +export LC_ALL=C +tab="$(printf '\t')" +join -t "$tab" \ + <(find papers -name title.md \ + -exec awk -v OFS='\t' '{ split(FILENAME, p, "/"); print p[1] "/" p[2], $0 }' {} + \ + | sort) \ + <(find papers -name abstract.md -printf '%h\t%s\n' | sort) +""" + +SENTENCE = ( + "We study the problem from first principles and report a method that " + "improves on the baseline across every setting we tried. " +) + + +def build_papers(root, lo, hi): + """Paper folders `lo` through `hi - 1`. Every seventh lacks a title and + every eleventh an abstract, so the join has to drop rows rather than + pair everything.""" + papers = root / "papers" + for i in range(lo, hi): + folder = papers / f"2609.{i:06d}" + folder.mkdir(parents=True) + if i % 7: + (folder / "title.md").write_text(f"Paper {i}: on the {i % 13}th method") + if i % 11: + (folder / "abstract.md").write_text(SENTENCE * (1 + i % 5)) + + +def native_rows(proc): + assert proc.returncode == 0, proc.stderr + rows = [] + for line in proc.stdout.splitlines(): + folder, title, size = line.split("\t") + rows.append((folder, title, int(size))) + return rows + + +def describe_join_speed_of_light(): + @pytest.fixture + def root(tmp_path): + tree = tmp_path / "arxiv" + try: + yield tree + finally: + shutil.rmtree(tree, ignore_errors=True) + + @pytest.mark.xfail( + strict=True, reason="path-table join is an unindexed nested loop" + ) + def it_matches_native_rows_within_the_bar(root): + startup = startup_seconds() + + def native(): + proc, seconds = timed(["bash", "-c", NATIVE], root, timeout=600) + return native_rows(proc), seconds + + n = grow_until_native_takes_a_second( + lambda lo, hi: build_papers(root, lo, hi), + native, + start=10_000, + ceiling=2**20, + ) + expected, native_seconds = best_of(native, "native") + assert len(expected) == sum(1 for i in range(n) if i % 7 and i % 11) + + def dirsql(): + proc, seconds = timed( + [cli(), "query", QUERY], + root, + timeout=dirsql_timeout(native_seconds), + ) + return dirsql_rows(proc, COLUMNS), seconds + + actual, dirsql_seconds = best_of( + dirsql, "dirsql", hopeless_seconds(native_seconds, startup) + ) + + assert actual == expected + assert_speed_of_light("join", native_seconds, dirsql_seconds, startup) diff --git a/packages/python/tests/e2e/use_case_notebook_test.py b/packages/python/tests/e2e/use_case_notebook_test.py new file mode 100644 index 000000000..fa394480f --- /dev/null +++ b/packages/python/tests/e2e/use_case_notebook_test.py @@ -0,0 +1,127 @@ +"""Use case 3, the lab notebook: one file, many rows. + +One markdown file of dated `## YYYY-MM-DD` sections, split into rows by a +parser script. Native is the parser alone piped into a sort-and-head. dirsql +runs the same parser once through `--on-file` and sorts in SQLite. No mocks: +real console script, real process, real filesystem, real parser spawn. +""" + +from __future__ import annotations + +import datetime +import shutil + +import pytest + +from .speed_of_light import ( + assert_speed_of_light, + best_of, + cli, + dirsql_rows, + dirsql_timeout, + grow_until_native_takes_a_second, + hopeless_seconds, + startup_seconds, + timed, +) + +COLUMNS = ("date", "words") +QUERY = "SELECT date, words FROM './notebook.md' ORDER BY date DESC LIMIT 5" + +EXTRACT = """\ +import json, re, sys + +text = open(sys.argv[1], encoding="utf-8").read() +rows = [ + {"date": date, "body": body.strip(), "words": len(body.split())} + for date, body in re.findall(r"^## (\\d{4}-\\d\\d-\\d\\d)\\n(.*?)(?=^## |\\Z)", text, re.S | re.M) +] +print(json.dumps(rows)) +""" + +TOP_FIVE = ( + "import json, sys;" + " rows = json.load(sys.stdin);" + " rows.sort(key=lambda r: r['date'], reverse=True);" + " [print(r['date'], r['words'], sep='\\t') for r in rows[:5]]" +) + +NATIVE = f'python3 extract.py notebook.md | python3 -c "{TOP_FIVE}"' + +WORDS = [ + "ran", + "the", + "cells", + "again", + "and", + "logged", + "every", + "reading", + "before", + "lunch", +] + + +def build_notebook(root, n): + """`n` sections, newest first, one per day back from a fixed date, each + with a distinct word count so the rows carry something to compare.""" + root.mkdir(parents=True, exist_ok=True) + (root / "extract.py").write_text(EXTRACT) + newest = datetime.date(2026, 9, 23) + chunks = ["# Lab notebook\n\nOne file. Newest entry at the top.\n"] + for i in range(n): + day = newest - datetime.timedelta(days=i) + body = " ".join(WORDS[j % len(WORDS)] for j in range(3 + i % 37)) + chunks.append(f"\n## {day.isoformat()}\n\n{body}\n") + (root / "notebook.md").write_text("".join(chunks)) + + +def native_rows(proc): + assert proc.returncode == 0, proc.stderr + rows = [] + for line in proc.stdout.splitlines(): + date, words = line.split("\t") + rows.append((date, int(words))) + return rows + + +def describe_notebook_speed_of_light(): + @pytest.fixture + def root(tmp_path): + tree = tmp_path / "notebook" + try: + yield tree + finally: + shutil.rmtree(tree, ignore_errors=True) + + @pytest.mark.xfail( + strict=True, + reason="ingesting the parser's JSON rows is slower than a second interpreter's json.load and sort", + ) + def it_matches_native_rows_within_the_bar(root): + startup = startup_seconds() + + def native(): + proc, seconds = timed(["sh", "-c", NATIVE], root, timeout=600) + return native_rows(proc), seconds + + grow_until_native_takes_a_second( + lambda _, n: build_notebook(root, n), native, start=20_000, ceiling=2**22 + ) + expected, native_seconds = best_of(native, "native") + assert len(expected) == 5, expected + + def dirsql(): + proc, seconds = timed( + [cli(), "query", QUERY, "--on-file", "python3 extract.py {path}"], + root, + timeout=dirsql_timeout(native_seconds), + ) + return dirsql_rows(proc, COLUMNS), seconds + + actual, dirsql_seconds = best_of( + dirsql, "dirsql", hopeless_seconds(native_seconds, startup) + ) + + assert actual == expected + assert_speed_of_light("notebook", native_seconds, dirsql_seconds, startup) diff --git a/packages/python/tests/e2e/use_case_transcripts_test.py b/packages/python/tests/e2e/use_case_transcripts_test.py new file mode 100644 index 000000000..21058d5b2 --- /dev/null +++ b/packages/python/tests/e2e/use_case_transcripts_test.py @@ -0,0 +1,200 @@ +"""Use case 4, the transcripts: many files, many rows. + +A tree of JSONL session logs, each turned into message rows by a parser +named in `.dirsql.toml`; the question is how many messages each project has +per role. Native is one parser process over every file. dirsql spawns the +same parser per file through the real launcher. No mocks: real console +script, real process, real filesystem, real parser spawns. +""" + +from __future__ import annotations + +import json +import shutil + +import pytest + +from .speed_of_light import ( + assert_speed_of_light, + best_of, + cli, + dirsql_rows, + dirsql_timeout, + grow_until_native_takes_a_second, + hopeless_seconds, + startup_seconds, + timed, +) + +COLUMNS = ("project", "role", "n") +QUERY = ( + "SELECT project, role, count(*) AS n FROM messages" + " GROUP BY project, role ORDER BY project, role" +) + +CONFIG = """\ +[[table]] +name = "messages" +glob = "projects/**/*.jsonl" +ddl = "CREATE TABLE messages (project TEXT, kind TEXT, ts TEXT, role TEXT, text TEXT)" +on-file = "python3 messages.py {path}" +""" + +MESSAGES = """\ +import json, sys + +NOISE = ("<", "Another Claude session", "Caveat:", "This session is being continued") + + +def rows(path): + parts = path.split("/") + project = parts[parts.index("projects") + 1] + kind = "subagent" if "subagents" in parts else "main" + out = [] + for line in open(path, encoding="utf-8"): + if not line.strip(): + continue + try: + entry = json.loads(line) + except ValueError: + continue + if entry.get("type") not in ("user", "assistant"): + continue + content = (entry.get("message") or {}).get("content") + if isinstance(content, str): + text = content + else: + text = "\\n".join(b["text"] for b in (content or []) if b.get("type") == "text") + if not text or text.startswith(NOISE): + continue + out.append({"project": project, "kind": kind, "ts": entry.get("timestamp"), "role": entry["type"], "text": text}) + return out + + +if __name__ == "__main__": + print(json.dumps(rows(sys.argv[1]))) +""" + +COUNT = """\ +import os, sys +from collections import Counter +from messages import rows + +counts = Counter() +for folder, _, names in os.walk("projects"): + for name in names: + if name.endswith(".jsonl"): + for row in rows(os.path.join(folder, name)): + counts[(row["project"], row["role"])] += 1 +for (project, role), n in sorted(counts.items()): + print(project, role, n, sep="\\t") +""" + +NATIVE = "python3 count.py" + +LINE_KINDS = ( + { + "type": "user", + "message": {"content": "please run the tests again and show me the output"}, + }, + { + "type": "assistant", + "message": { + "content": [ + {"type": "text", "text": "Running them now."}, + {"type": "tool_use", "name": "Bash"}, + ] + }, + }, + { + "type": "assistant", + "message": {"content": [{"type": "tool_use", "name": "Read"}]}, + }, + { + "type": "user", + "message": {"content": "ignored"}, + }, + {"type": "progress", "data": {"step": 3}}, + { + "type": "assistant", + "message": {"content": [{"type": "text", "text": "Caveat: this is noise"}]}, + }, +) +PADDING = " ".join(["filler"] * 60) + + +def build_transcripts(root, lo, hi): + """Session files `lo` through `hi - 1`, two hundred lines each across + eight projects, one in four under a `subagents/` folder, each line a few + hundred bytes.""" + if lo == 0: + root.mkdir(parents=True) + (root / ".dirsql.toml").write_text(CONFIG) + (root / "messages.py").write_text(MESSAGES) + (root / "count.py").write_text(COUNT) + for i in range(lo, hi): + project = root / "projects" / f"project-{i % 8}" + if i % 4 == 0: + path = project / f"session-{i}" / "subagents" / "agent.jsonl" + else: + path = project / f"session-{i}.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + lines = [] + for j in range(200): + entry = dict(LINE_KINDS[(i + j) % len(LINE_KINDS)]) + entry["timestamp"] = f"2026-09-{1 + j % 28:02d}T10:{j % 60:02d}:00Z" + entry["uuid"] = f"{i}-{j}" + entry["padding"] = PADDING + lines.append(json.dumps(entry)) + path.write_text("\n".join(lines) + "\n") + + +def native_rows(proc): + assert proc.returncode == 0, proc.stderr + rows = [] + for line in proc.stdout.splitlines(): + project, role, n = line.split("\t") + rows.append((project, role, int(n))) + return rows + + +def describe_transcripts_speed_of_light(): + @pytest.fixture + def root(tmp_path): + tree = tmp_path / "transcripts" + try: + yield tree + finally: + shutil.rmtree(tree, ignore_errors=True) + + @pytest.mark.xfail(strict=True, reason="on-file spawns one parser process per file") + def it_matches_native_rows_within_the_bar(root): + startup = startup_seconds() + + def native(): + proc, seconds = timed(["sh", "-c", NATIVE], root, timeout=600) + return native_rows(proc), seconds + + grow_until_native_takes_a_second( + lambda lo, hi: build_transcripts(root, lo, hi), + native, + start=128, + ceiling=2**16, + ) + expected, native_seconds = best_of(native, "native") + assert len(expected) == 16, expected + + def dirsql(): + proc, seconds = timed( + [cli(), "query", QUERY, "-c", ".dirsql.toml"], + root, + timeout=dirsql_timeout(native_seconds), + ) + return dirsql_rows(proc, COLUMNS), seconds + + actual, dirsql_seconds = best_of( + dirsql, "dirsql", hopeless_seconds(native_seconds, startup) + ) + + assert actual == expected + assert_speed_of_light("transcripts", native_seconds, dirsql_seconds, startup) diff --git a/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light.py b/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light.py new file mode 100644 index 000000000..66173c7f1 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light.py @@ -0,0 +1,124 @@ +"""Harness for the use-case speed-of-light tests. + +Each use-case test asks one question of a generated corpus twice: once through +the fastest native pipeline (find, sort, join, awk, a single parser process) +and once through the real `dirsql` launcher. The rows must be identical and +dirsql's wall time, after subtracting the launcher's measured startup, must be +within ten percent of native. The bar is a ratio, never a millisecond number, +so it holds on any machine. + +Startup is the mean of ten `dirsql --help` runs; `--help` and `SELECT 1` cost +the same, so it stands in for the fixed per-invocation cost. A corpus too +small to time is grown by doubling until the native pipeline takes at least a +second, so the comparison is above the noise floor. Each side is run three +times and the fastest run counts. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import time + +ONE_SECOND = 1.0 +TOLERANCE = 1.1 +STARTUP_PROBES = 10 +BEST_OF = 3 + + +def cli() -> str: + dirsql = shutil.which("dirsql") + assert dirsql is not None, ( + "`dirsql` console script not on PATH -- run `uv run maturin develop`" + ) + return dirsql + + +def timed(argv, cwd, timeout, env=None): + started = time.perf_counter() + proc = subprocess.run( + argv, + cwd=str(cwd), + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=timeout, + env=env, + check=False, + ) + return proc, time.perf_counter() - started + + +def startup_seconds(argv=None, cwd=None) -> float: + argv = [cli(), "--help"] if argv is None else argv + samples = [] + for _ in range(STARTUP_PROBES): + proc, seconds = timed(argv, cwd or ".", timeout=60) + assert proc.returncode == 0, proc.stderr + samples.append(seconds) + return sum(samples) / len(samples) + + +def dirsql_rows(proc, columns): + assert proc.returncode == 0, ( + f"dirsql failed: stdout={proc.stdout!r} stderr={proc.stderr!r}" + ) + return [tuple(row[column] for column in columns) for row in json.loads(proc.stdout)] + + +def grow_until_native_takes_a_second(build, native, start, ceiling): + """Grow the corpus by doubling until `native` runs for a second. + + `build(lo, hi)` adds items `lo` through `hi - 1` to the corpus already + holding `lo` items, so each doubling writes only the new half. + """ + built, size = 0, start + while True: + build(built, size) + built = size + _, seconds = native() + if seconds >= ONE_SECOND or size >= ceiling: + return size + size *= 2 + + +def best_of(run, name, hopeless=None): + """Run `run` three times; return the rows and wall time of the fastest. + + `run` returns `(rows, seconds)`. A run that raises `TimeoutExpired` fails + the test immediately with the timeout it blew through. A run slower than + `hopeless` seconds ends the series early: noise does not account for an + order of magnitude, so no retry could bring the best under the bar. + """ + fastest = None + for _ in range(BEST_OF): + try: + rows, seconds = run() + except subprocess.TimeoutExpired as expired: + raise AssertionError( + f"{name} did not finish within {expired.timeout:.0f}s" + ) from None + if fastest is None or seconds < fastest[1]: + fastest = (rows, seconds) + if hopeless is not None and seconds > hopeless: + break + return fastest + + +def hopeless_seconds(native_seconds: float, startup: float) -> float: + return startup + 10 * TOLERANCE * native_seconds + + +def dirsql_timeout(native_seconds: float) -> float: + return max(60.0, 30 * native_seconds) + + +def assert_speed_of_light(name, native_seconds, dirsql_seconds, startup): + measured = dirsql_seconds - startup + budget = TOLERANCE * native_seconds + assert measured <= budget, ( + f"{name}: dirsql {dirsql_seconds:.3f}s minus startup {startup:.3f}s" + f" = {measured:.3f}s, over {TOLERANCE}x native {native_seconds:.3f}s" + f" = {budget:.3f}s" + ) diff --git a/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light_test.py b/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light_test.py new file mode 100644 index 000000000..5fce8cd4f --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light_test.py @@ -0,0 +1,131 @@ +"""Use case 2 through the plugin: embedding every abstract in a corpus. + +The native reference is the model itself: one Python process that loads the +same model2vec model and batch-encodes every file's text. dirsql embeds each +file through `embed()` and the real worker process, cold cache every run. +The vectors must be identical and dirsql's wall time, after subtracting the +launcher's startup, within ten percent of the batch encode. Nothing mocked: +real launcher, real worker spawn, real model2vec inference on the on-disk +model from conftest. +""" + +from __future__ import annotations + +import json +import os +import shutil +import sys +from importlib import resources + +import pytest + +from .speed_of_light import ( + assert_speed_of_light, + best_of, + dirsql_rows, + dirsql_timeout, + grow_until_native_takes_a_second, + hopeless_seconds, + startup_seconds, + timed, +) + +_FRAGMENT = str(resources.files("dirsql_plugin_embeddings").joinpath("dirsql.toml")) +LAUNCHER = [ + sys.executable, + "-c", + "import sys; from dirsql.cli.main import main; sys.exit(main())", + "--no-plugin", +] +COLUMNS = ("path", "emb") + +ENCODE = """\ +import json, os, sys +from model2vec import StaticModel + +model = StaticModel.from_pretrained(sys.argv[1]) +names = sorted(os.listdir("abstracts")) +texts = [open(os.path.join("abstracts", name), encoding="utf-8").read() for name in names] +for name, vector in zip(names, model.encode(texts, show_progress_bar=False)): + print(name, json.dumps([float(c) for c in vector]), sep="\\t") +""" + +VOCABULARY = ("hello", "world") + + +def build_abstracts(root, lo, hi): + """Abstracts `lo` through `hi - 1`, each a different run of the model's + two known words so no two neighbours embed alike.""" + folder = root / "abstracts" + if lo == 0: + folder.mkdir(parents=True) + (root / "encode.py").write_text(ENCODE) + for i in range(lo, hi): + words = [VOCABULARY[(i >> bit) & 1] for bit in range(3 + i % 9)] + (folder / f"{i:07d}.md").write_text(" ".join(words), encoding="utf-8") + + +def native_rows(proc): + assert proc.returncode == 0, proc.stderr + rows = [] + for line in proc.stdout.splitlines(): + name, vector = line.split("\t") + rows.append((name, json.loads(vector))) + return rows + + +def describe_embeddings_speed_of_light(): + @pytest.fixture + def root(tmp_path): + tree = tmp_path / "corpus" + try: + yield tree + finally: + shutil.rmtree(tree, ignore_errors=True) + + @pytest.mark.xfail( + strict=True, + reason="embed() makes one worker round trip per value instead of a batch", + ) + def it_matches_the_batch_encode_within_the_bar(root, tiny_model, tmp_path): + startup = startup_seconds([*LAUNCHER, "--help"], cwd=tmp_path) + + def native(): + proc, seconds = timed( + [sys.executable, "encode.py", tiny_model], root, timeout=600 + ) + return native_rows(proc), seconds + + grow_until_native_takes_a_second( + lambda lo, hi: build_abstracts(root, lo, hi), + native, + start=4_096, + ceiling=2**18, + ) + expected, native_seconds = best_of(native, "native") + + query = ( + "SELECT path, embed(content, '" + tiny_model + "') AS emb" + " FROM './abstracts/*.md' ORDER BY path" + ) + cold_runs = iter(range(10)) + + def dirsql(): + cache_home = tmp_path / f"xdg-cache-{next(cold_runs)}" + proc, seconds = timed( + [*LAUNCHER, "query", query, "--config", _FRAGMENT], + root, + timeout=dirsql_timeout(native_seconds), + env={**os.environ, "XDG_CACHE_HOME": str(cache_home)}, + ) + rows = dirsql_rows(proc, COLUMNS) + return [ + (os.path.basename(path), json.loads(emb)) for path, emb in rows + ], seconds + + actual, dirsql_seconds = best_of( + dirsql, "dirsql", hopeless_seconds(native_seconds, startup) + ) + + assert actual == expected + assert_speed_of_light("embeddings", native_seconds, dirsql_seconds, startup)