github-featured/omnigent-ai--omnigent
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

perf(benchmark): bulk-insert the SQLite seed corpus in one transaction (#2947)

Browse Source 
dev/benchmarks/omnigent/seed.py seeded the benchmark corpus through the
production store ORM API one row at a time (~2M single-row INSERTs, ~20k
commits, each preceded by throwaway PRAGMAs on session open), taking ~6-10
min on CI. The benchmark only measures the store read path, so the write
strategy does not taint what's measured provided the resulting corpus is the
same shape.

Add a SQLAlchemy Core bulk-insert fast path (_seed_via_core) that writes the
whole corpus in one transaction via ~10 batched executemany flushes (1 commit
instead of ~20k). It uses the ORM Table objects so Uuid16 binds bare-hex to
16 bytes byte-identically to the store, computes title_hash explicitly
(Python defaults don't fire under executemany, and sets all kind/status
columns explicitly. The schema at head carries no FK constraints (migration
p1a2b3c4d5e6 dropped them all), so insert order is free under
PRAGMA foreign_keys=ON.

Dialect-gated: SQLite uses the fast path; every other dialect (e.g. the
nightly Postgres benchmark) falls back to the existing store-API loop
(_seed_via_store), extracted verbatim, so behavior there stays identical.

Byte-stable: same RNG seed/counts/_FRAGMENTS, same generate_*_id calls, same
per-session draw order (title first, then items), same 0-based position
allocation, same label stamped on the last session, same _meta_value config
string. Item data/search_text are built byte-identical to
MessageData.model_dump(exclude_none=True) + extract_search_text (the slow
path keeps _make_items as the single source of truth). The fast path item
build bypasses pydantic (building plain dicts) to keep the 1M-item Python
phase cheap; a byte-stability test pins both paths to identical corpora.

Idempotency preserved: the reuse-skip check, --reseed, and --print-head work
unchanged; ensure_user(local) and the seed-meta label upsert are mirrored
via sqlite_insert.on_conflict_do_*.

Target: ~20-30s end-to-end (was ~6-10 min) for the 5000x200 corpus; measured
~27s locally. Scope: seed.py + a new test file only; no product store/db code
under omnigent/stores/ or omnigent/db/ touched.
EOF
)
This commit is contained in:
Zeyi (Rice) Fan
2026-07-20 18:46:34 -07:00
committed by GitHub
parent 943bafe342
commit dc64952419
2 changed files with 517 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dev/benchmarks/omnigent/seed.py
+313 -8
View File
@@ -29,6 +29,8 @@ Run standalone::
from __future__ import annotations
import argparse
import hashlib
import json
import random
import sys
from pathlib import Path
@@ -36,7 +38,37 @@ from pathlib import Path
# Allow ``uv run <path>`` (no package context) to import omnigent + siblings.
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
from omnigent.db.utils import _get_head_db_revision, generate_agent_id
from sqlalchemy import text
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
from sqlalchemy.engine.url import make_url
from omnigent.db.db_models import (
LABEL_VALUE_MAX_LEN,
SqlAgent,
SqlConversation,
SqlConversationItem,
SqlConversationLabel,
SqlConversationMetadata,
SqlSessionPermission,
SqlUser,
current_workspace_id,
)
from omnigent.db.enum_codecs import (
encode_agent_kind,
encode_conversation_kind,
encode_item_status,
encode_item_type,
)
from omnigent.db.utils import (
_FTS_TABLE,
_get_head_db_revision,
generate_agent_id,
generate_conversation_id,
generate_item_id,
get_or_create_engine,
now_epoch,
strip_nul_bytes,
)
from omnigent.entities import MessageData, NewConversationItem
from omnigent.server.auth import LEVEL_OWNER, RESERVED_USER_LOCAL
from omnigent.stores.conversation_store.sqlalchemy_store import SqlAlchemyConversationStore
@@ -48,6 +80,7 @@ _SEED_META_LABEL = "omni_bench_seed"
# Fixed identifiers so the corpus is byte-stable across runs at a given config.
_AGENT_NAME = "bench-agent"
_AGENT_BUNDLE = "bench/seed" # never validated on the read path
_DEFAULT_SESSIONS = 5000
_DEFAULT_ITEMS = 50
_DEFAULT_RNG_SEED = 1234
@@ -67,6 +100,19 @@ _FRAGMENTS = (
"reproduce the elicitation race on reconnect",
)
# FTS5 mirror row written per item on SQLite (must match omnigent.db.utils
# ``insert_fts`` / ``_FTS_TABLE``). Bound by name in :data:`_FTS_INSERT_SQL`.
_FTS_INSERT_SQL = text(
f"INSERT INTO {_FTS_TABLE} (item_id, conversation_id, search_text) "
"VALUES (:item_id, :cid, :st)"
)
# Rows buffered per Core ``executemany`` flush. Only the item/FTS buffers
# (1:1 with items) approach this; the per-session tables are held in full
# (a few thousand rows) and inserted in one shot each. 100k keeps a 5000×200
# corpus to ~10 flushes per table and bounds peak memory to a few tens of MB.
_CORE_ITEM_CHUNK = 100_000
def _meta_value(sessions: int, items_per_session: int, rng_seed: int, head: str) -> str:
"""Serialize the corpus config into the seed-marker label value.
@@ -99,17 +145,23 @@ def _make_items(rng: random.Random, count: int) -> list[NewConversationItem]:
"""
items: list[NewConversationItem] = []
for i in range(count):
text = f"{rng.choice(_FRAGMENTS)} (item {i})"
text_str = f"{rng.choice(_FRAGMENTS)} (item {i})"
items.append(
NewConversationItem(
type="message",
response_id=f"resp_seed_{i}",
data=MessageData(role="user", content=[{"type": "input_text", "text": text}]),
data=MessageData(role="user", content=[{"type": "input_text", "text": text_str}]),
)
)
return items
def _progress(s: int, sessions: int) -> None:
"""Print a coarse progress line every 10% (only for sizeable corpora)."""
if sessions >= 100 and s % (sessions // 10) == 0 and s:
print(f"seed: {s}/{sessions} sessions")
def seed(
db_uri: str,
*,
@@ -117,6 +169,7 @@ def seed(
items_per_session: int = _DEFAULT_ITEMS,
rng_seed: int = _DEFAULT_RNG_SEED,
reseed: bool = False,
_fast: bool | None = None,
) -> int:
"""Seed *sessions* sessions × *items_per_session* items into *db_uri*.
@@ -131,10 +184,21 @@ def seed(
:param items_per_session: Conversation items appended to each session.
:param rng_seed: Seed for the deterministic text RNG.
:param reseed: Seed even when a matching corpus is already present.
:param _fast: Override the write strategy. ``None`` (default) uses the
bulk-insert Core fast path for SQLite and the store-API loop for every
other dialect; ``True`` forces the fast path (falls back to the loop on
non-SQLite); ``False`` forces the store-API loop (used by the
byte-stability test to compare both paths on SQLite).
:returns: The number of sessions created (0 when a matching seed is reused).
"""
conv = SqlAlchemyConversationStore(db_uri)
perms = SqlAlchemyPermissionStore(db_uri)
dialect = make_url(db_uri).get_backend_name()
use_fast = (dialect == "sqlite") if _fast is None else bool(_fast)
if use_fast and dialect != "sqlite":
# The fast path is SQLite-only (FTS5 + single-transaction bulk insert);
# a forced fast request on another dialect degrades to the store loop.
use_fast = False
# Read the current schema head at runtime (no DB contacted) and fold it into
# the reuse marker, so a corpus from an older schema is auto-reseeded.
@@ -149,6 +213,44 @@ def seed(
print(f"seed: existing corpus differs ({existing!r} != {want!r}); pass --reseed")
return 0
if use_fast:
n = _seed_via_core(
db_uri,
sessions=sessions,
items_per_session=items_per_session,
rng_seed=rng_seed,
want=want,
)
else:
perms = SqlAlchemyPermissionStore(db_uri)
n = _seed_via_store(
conv,
perms,
sessions=sessions,
items_per_session=items_per_session,
rng_seed=rng_seed,
want=want,
)
print(f"seed: created {n} sessions × {items_per_session} items ({want})")
return n
def _seed_via_store(
conv: SqlAlchemyConversationStore,
perms: SqlAlchemyPermissionStore,
*,
sessions: int,
items_per_session: int,
rng_seed: int,
want: str,
) -> int:
"""Seed through the production store ORM API (one row/commit at a time).
This is the original path and the only one used on non-SQLite dialects
(e.g. the nightly Postgres benchmark). It is kept verbatim so behavior
there stays identical.
"""
perms.ensure_user(RESERVED_USER_LOCAL)
rng = random.Random(rng_seed)
@@ -157,7 +259,7 @@ def seed(
created = conv.create_session_with_agent(
agent_id=generate_agent_id(),
agent_name=_AGENT_NAME,
agent_bundle_location="bench/seed", # never validated on the read path
agent_bundle_location=_AGENT_BUNDLE,
agent_description=None,
title=f"bench session {s}: {rng.choice(_FRAGMENTS)}",
)
@@ -166,8 +268,7 @@ def seed(
perms.grant(RESERVED_USER_LOCAL, sid, LEVEL_OWNER)
if items_per_session:
conv.append(sid, _make_items(rng, items_per_session))
if sessions >= 100 and s % (sessions // 10) == 0 and s:
print(f"seed: {s}/{sessions} sessions")
_progress(s, sessions)
# Stamp the corpus config on the LAST (newest) session — that's the one
# ``_existing_seed_meta``'s default desc listing returns, so the reuse
@@ -175,7 +276,211 @@ def seed(
if last_sid:
conv.set_labels(last_sid, {_SEED_META_LABEL: want})
print(f"seed: created {sessions} sessions × {items_per_session} items ({want})")
return sessions
def _seed_via_core(
db_uri: str,
*,
sessions: int,
items_per_session: int,
rng_seed: int,
want: str,
) -> int:
"""Seed the entire corpus in one transaction via SQLAlchemy Core.
Writes the same DB rows the store-API loop would, but batches them into a
handful of ``executemany`` flushes under a single ``BEGIN``/``COMMIT`` —
~10 batched INSERTs and 1 commit instead of ~2M single-row INSERTs and
~20k commits. The schema at head carries no FK constraints (migration
``p1a2b3c4d5e6`` dropped them all), so insert order is free and the
engine's ``PRAGMA foreign_keys=ON`` enforces nothing.
The RNG draw order and the per-row serialization are kept identical to the
store path: per session, the title fragment is drawn first, then the
``items_per_session`` item fragments. Item ``data`` is
``strip_nul_bytes(json.dumps(...))`` (default separators) of a plain dict
that mirrors ``MessageData.model_dump(exclude_none=True)``, and
``search_text`` mirrors ``extract_search_text``'s message branch — both
built directly (no pydantic) so the 1M-item Python build stays cheap. So a
corpus seeded here is the same shape (same ids-space, same text, same
positions, same labels) as one seeded through the store — only the write
strategy differs. ``tests/benchmarks/test_seed_fast_path.py`` pins the two
paths to identical corpora.
"""
engine = get_or_create_engine(db_uri)
ws = current_workspace_id()
rng = random.Random(rng_seed)
# Per-session scalar rows (conversations/agents/metadata/permissions) are
# small (a few thousand); hold them in full and insert each in one shot.
conv_rows: list[dict] = []
agent_rows: list[dict] = []
meta_rows: list[dict] = []
perm_rows: list[dict] = []
# Items + FTS mirror are 1:1 with items and dominate the volume (1M+ for a
# full seed); stream them in chunks to bound memory while staying in the
# single transaction.
item_buf: list[dict] = []
fts_buf: list[dict] = []
last_sid = ""
with engine.begin() as conn:
# ensure_user("local") — ON CONFLICT DO NOTHING, mirroring the store.
conn.execute(
sqlite_insert(SqlUser)
.values(workspace_id=ws, id=RESERVED_USER_LOCAL, is_admin=False)
.on_conflict_do_nothing(index_elements=["workspace_id", "id"])
)
for s in range(sessions):
now = now_epoch()
agent_id = generate_agent_id()
sid = generate_conversation_id()
# RNG draw order matches the store path: title first, then items.
title = f"bench session {s}: {rng.choice(_FRAGMENTS)}"
last_sid = sid
conv_rows.append(
{
"workspace_id": ws,
"id": sid,
"created_at": now,
"updated_at": now,
"title": title,
"title_hash": hashlib.sha256(title.encode("utf-8")).digest()[:16],
"parent_conversation_id": None,
"root_conversation_id": sid,
"next_position": items_per_session,
"agent_id": agent_id,
"session_overrides": None,
"archived": False,
}
)
agent_rows.append(
{
"workspace_id": ws,
"id": agent_id,
"created_at": now,
"name": _AGENT_NAME,
"bundle_location": _AGENT_BUNDLE,
"version": 1,
"kind": encode_agent_kind("session"),
"description": None,
"updated_at": None,
}
)
meta_rows.append(
{
"workspace_id": ws,
"id": sid,
"kind": encode_conversation_kind("default"),
"runner_id": None,
"host_id": None,
"sub_agent_name": None,
"external_session_id": None,
"session_state": None,
"session_usage": None,
"terminal_launch_args": None,
"workspace": None,
"git_branch": None,
"runner_last_seen": None,
"live_status": None,
"pending_elicitation_count": None,
}
)
perm_rows.append(
{
"workspace_id": ws,
"user_id": RESERVED_USER_LOCAL,
"conversation_id": sid,
"level": LEVEL_OWNER,
}
)
if items_per_session:
# Build item payloads straight to dicts (no pydantic) so the
# 1M-item Python build stays cheap. The output is byte-identical
# to the store path: the text format mirrors ``_make_items``,
# the ``data`` dict mirrors ``MessageData.model_dump(exclude_none=
# True)``, and ``search`` mirrors ``extract_search_text``'s
# message branch. The byte-stability test
# (tests/benchmarks/test_seed_fast_path.py) pins this to the
# store path's rows. ``_make_items`` is still used by the slow
# path above, so the text format stays single-sourced there.
for i in range(items_per_session):
text_str = f"{rng.choice(_FRAGMENTS)} (item {i})"
data_dict = {
"role": "user",
"content": [{"type": "input_text", "text": text_str}],
}
data = strip_nul_bytes(json.dumps(data_dict))
search = strip_nul_bytes(
" ".join(
block["text"]
for block in data_dict["content"]
if isinstance(block, dict) and block.get("text")
)
)
item_id = generate_item_id("message")
item_buf.append(
{
"workspace_id": ws,
"conversation_id": sid,
"id": item_id,
"response_id": f"resp_seed_{i}",
"created_at": now,
"status": encode_item_status("completed"),
"position": i,
"type": encode_item_type("message"),
"data": data,
"search_text": search,
"created_by": None,
}
)
fts_buf.append({"item_id": item_id, "cid": sid, "st": search})
if len(item_buf) >= _CORE_ITEM_CHUNK:
conn.execute(SqlConversationItem.__table__.insert(), item_buf)
conn.execute(_FTS_INSERT_SQL, fts_buf)
item_buf.clear()
fts_buf.clear()
_progress(s, sessions)
if item_buf:
conn.execute(SqlConversationItem.__table__.insert(), item_buf)
conn.execute(_FTS_INSERT_SQL, fts_buf)
item_buf.clear()
fts_buf.clear()
# No FKs at head → order is free; insert the per-session scalar tables
# now (after the streamed items) in one shot each.
conn.execute(SqlConversation.__table__.insert(), conv_rows)
conn.execute(SqlAgent.__table__.insert(), agent_rows)
conn.execute(SqlConversationMetadata.__table__.insert(), meta_rows)
conn.execute(SqlSessionPermission.__table__.insert(), perm_rows)
# Stamp the corpus config on the LAST (newest) session, matching the
# store path's ``set_labels`` upsert (clamped to LABEL_VALUE_MAX_LEN).
if last_sid:
label_now = now_epoch()
label_value = want[:LABEL_VALUE_MAX_LEN]
conn.execute(
sqlite_insert(SqlConversationLabel)
.values(
workspace_id=ws,
conversation_id=last_sid,
key=_SEED_META_LABEL,
value=label_value,
updated_at=label_now,
)
.on_conflict_do_update(
index_elements=["workspace_id", "conversation_id", "key"],
set_={"value": label_value, "updated_at": label_now},
)
)
return sessions
tests/benchmarks/test_seed_fast_path.py
+204
View File
@@ -0,0 +1,204 @@
"""Tests for the benchmark corpus seeder's bulk-insert fast path.
The seeder has two write strategies selected by dialect:
* the production store-ORM loop (one row/commit at a time) — the only path on
non-SQLite dialects such as the nightly Postgres benchmark; and
* a SQLAlchemy-Core bulk-insert fast path (one transaction, ~10 batched
``executemany`` flushes) — SQLite only.
These tests pin the fast path to the store path's contract: the same corpus
*shape* (row counts, RNG-determined text, JSON ``data`` blobs, position
allocation, labels) so the read journeys measure the same volume either way.
"""
from __future__ import annotations
import os
from pathlib import Path
import pytest
from sqlalchemy import text
from sqlalchemy.engine.url import make_url
from omnigent.db.utils import get_or_create_engine
from omnigent.server.auth import LEVEL_OWNER, RESERVED_USER_LOCAL
from omnigent.stores.conversation_store.sqlalchemy_store import SqlAlchemyConversationStore
# ── helpers ──────────────────────────────────────────────────
def _count(engine, table: str) -> int:
with engine.connect() as conn:
return conn.execute(text(f"SELECT COUNT(*) FROM {table}")).scalar_one()
def _s_of(title: str) -> int:
"""Recover the session index from a seeded title ``"bench session {s}: ..."``."""
return int(title.split(":", 1)[0][len("bench session ") :])
def _titles_ordered(engine):
"""Titles keyed by session index s, sorted (RNG-determined, order-sensitive)."""
with engine.connect() as conn:
rows = conn.execute(text("SELECT title FROM conversations")).all()
return sorted((_s_of(t), t) for (t,) in rows)
def _item_text_ordered(engine):
"""(session s, position, search_text) sorted — the RNG draw sequence in order."""
with engine.connect() as conn:
convs = conn.execute(text("SELECT id, title FROM conversations")).all()
s_by_cid = {cid: _s_of(t) for cid, t in convs}
items = conn.execute(
text("SELECT conversation_id, position, search_text FROM conversation_items")
).all()
return sorted((s_by_cid[cid], pos, st) for cid, pos, st in items)
def _item_data_ordered(engine):
"""(session s, position, data-JSON) sorted — proves the serialization matches."""
with engine.connect() as conn:
convs = conn.execute(text("SELECT id, title FROM conversations")).all()
s_by_cid = {cid: _s_of(t) for cid, t in convs}
items = conn.execute(
text("SELECT conversation_id, position, data FROM conversation_items")
).all()
return sorted((s_by_cid[cid], pos, d) for cid, pos, d in items)
# ── fast path: row counts + read path through the store API ──
def test_seed_fast_path_row_counts_and_read_path(tmp_path: Path) -> None:
"""The fast path writes every table the store path would, listable as "local"."""
from dev.benchmarks.omnigent import seed as seed_mod
db_uri = f"sqlite:///{tmp_path / 'fast.db'}"
created = seed_mod.seed(db_uri, sessions=50, items_per_session=10, _fast=True)
assert created == 50
engine = get_or_create_engine(db_uri)
assert _count(engine, "conversations") == 50
assert _count(engine, "conversation_items") == 500
assert _count(engine, "conversation_items_fts") == 500
assert _count(engine, "agents") == 50
assert _count(engine, "omnigent_conversation_metadata") == 50
assert _count(engine, "session_permissions") == 50
assert _count(engine, "users") == 1
assert _count(engine, "conversation_labels") == 1
with engine.connect() as conn:
# next_position advanced to items_per_session on every conversation.
nps = conn.execute(text("SELECT next_position FROM conversations")).all()
assert all(r[0] == 10 for r in nps)
# every grant is LEVEL_OWNER for the reserved "local" user.
levels = conn.execute(
text("SELECT DISTINCT level FROM session_permissions WHERE user_id = :u"),
{"u": RESERVED_USER_LOCAL},
).all()
assert levels == [(LEVEL_OWNER,)]
# agent kind = session (2), metadata kind = default (1), archived = 0.
assert {r[0] for r in conn.execute(text("SELECT kind FROM agents")).all()} == {2}
assert {
r[0]
for r in conn.execute(text("SELECT kind FROM omnigent_conversation_metadata")).all()
} == {1}
assert {r[0] for r in conn.execute(text("SELECT archived FROM conversations")).all()} == {
0
}
# the seed-meta label is present and carries the corpus config.
labels = conn.execute(
text("SELECT value FROM conversation_labels WHERE key = :k"),
{"k": seed_mod._SEED_META_LABEL},
).all()
assert len(labels) == 1
assert labels[0][0].startswith("sessions=50;items=10;")
# Read path the benchmark measures (store API, not raw SQL).
conv = SqlAlchemyConversationStore(db_uri)
listing = conv.list_conversations(
limit=100,
agent_name="bench-agent",
accessible_by=RESERVED_USER_LOCAL,
has_agent_id=True,
)
assert len(listing.data) == 50 # all seeded sessions listable as "local"
items = conv.list_items(listing.data[0].id, limit=100)
assert len(items.data) == 10
# items come back in position order (response_id encodes the per-session index).
assert [i.response_id for i in items.data] == [f"resp_seed_{i}" for i in range(10)]
# Idempotent: a matching re-seed is a no-op (reuse-skip path).
assert seed_mod.seed(db_uri, sessions=50, items_per_session=10, _fast=True) == 0
# ── byte-stability: fast path replicates the store path's corpus ──
def test_seed_fast_path_corpus_matches_store_path(tmp_path: Path) -> None:
"""Same config + RNG seed → identical corpus shape via either write path.
The uuid4 ids are fresh per run, so we compare the RNG-determined content
(titles, per-session search_text in draw order, JSON ``data`` blobs) and the
row counts — proving the fast path draws the RNG in the same order and
serializes rows identically to the store ORM loop.
"""
from dev.benchmarks.omnigent import seed as seed_mod
cfg = {"sessions": 50, "items_per_session": 10, "rng_seed": 1234}
fast_uri = f"sqlite:///{tmp_path / 'fast.db'}"
slow_uri = f"sqlite:///{tmp_path / 'slow.db'}"
assert seed_mod.seed(fast_uri, _fast=True, **cfg) == 50
assert seed_mod.seed(slow_uri, _fast=False, **cfg) == 50 # force the store loop on SQLite
fe = get_or_create_engine(fast_uri)
se = get_or_create_engine(slow_uri)
for table in (
"conversations",
"conversation_items",
"conversation_items_fts",
"agents",
"omnigent_conversation_metadata",
"session_permissions",
"users",
"conversation_labels",
):
assert _count(fe, table) == _count(se, table), table
# RNG draw order + content match exactly (order-preserving).
assert _titles_ordered(fe) == _titles_ordered(se)
assert _item_text_ordered(fe) == _item_text_ordered(se)
# JSON serialization (default separators, exclude_none) matches byte-for-byte.
assert _item_data_ordered(fe) == _item_data_ordered(se)
# ── non-SQLite slow path (skipped unless a non-SQLite DB is provided) ──
def test_seed_slow_path_non_sqlite() -> None:
"""The store-API loop remains the path on non-SQLite dialects.
Skipped unless ``OMNIGENT_BENCH_NONSQLITE_URI`` points at a real non-SQLite
DB (e.g. the nightly Postgres benchmark). The fast path is SQLite-only, so
this guards the fallback the nightly run relies on.
"""
uri = os.environ.get("OMNIGENT_BENCH_NONSQLITE_URI")
if not uri:
pytest.skip(
"set OMNIGENT_BENCH_NONSQLITE_URI to a non-SQLite URI to exercise the slow seed path"
)
assert make_url(uri).get_backend_name() != "sqlite"
from dev.benchmarks.omnigent import seed as seed_mod
# --reseed so the count assertion holds against a DB that may already have a corpus.
created = seed_mod.seed(uri, sessions=20, items_per_session=5, reseed=True)
assert created == 20
conv = SqlAlchemyConversationStore(uri)
listing = conv.list_conversations(limit=100, agent_name="bench-agent")
assert len(listing.data) >= 20