Store context_snapshot compressed
One column is 89% of the database and the free tier allows 512 MB. Reads were already solved -- the column is deferred, so a page load never touches it and one screen fetches one row at a time -- but nothing had costed storage, and storage is the constraint with a cliff: 99.6 MB used, ~94 kB of disk per action, so the ceiling arrives around 5,400 actions and 944 are stored. Postgres already compresses it and only gets 1.7x. pglz is tuned for fast decompression of data a query might filter on, and nothing has ever filtered on an assembled prompt -- it is written once and read whole, rarely, by the Insights viewer. zlib gets 3.5x on the same text for a decompress on a request that already made an LLM call. Done as a TypeDecorator rather than a second column, so every call site still writes a dict and reads a dict back, and deferred/undefer/load_only keep naming the same attribute. Only the storage format moves. Migrations 43-45: add the bytea, convert into it, drop the original, rename. The backfill is the one destructive step in the file -- 44 removes the only other copy -- so it decompresses every row and compares it against what went in, and a row that fails aborts the run. The whole loop is one transaction, so an abort rolls the DROP back and the prompts are still there. Verified on real Postgres, replaying 43-45 from a pre-43 schema on a throwaway Neon database: 720,864 B of JSON became 204,293 B of bytea, 3.53x, the column came out named context_snapshot, every snapshot compared equal and the one NULL stayed NULL. Postgres does not return the disk by itself: DROP COLUMN only marks the column gone and the backfill leaves a dead tuple per row, so the table peaks near twice its size before settling. The deploy needs one VACUUM FULL to collect it; the migration comment says so. The egress fixture's snapshots are prose now rather than "x" * 20_000, and the prose generator moved to tools/fakeprose.py so the harness and the tests share one definition. A repeated character compresses a thousandfold: against the old fixture a compressed column looked free and the byte ceilings would have been guarding nothing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017Dvvqn9ZDR4ixeFPHNbww7
This commit is contained in:
parththakkar106
co-authored by
Claude Opus 5
Claude Opus 5
parent
a6cb49293c
commit
ae6e5af6c7
@@ -27,6 +27,7 @@ os.environ.pop("AIDND_DATABASE_URL", None)
|
||||
os.environ.pop("DATABASE_URL", None)
|
||||
|
||||
import json
|
||||
import random
|
||||
|
||||
import pytest
|
||||
from fastapi import Depends
|
||||
@@ -39,12 +40,20 @@ from app.context import history
|
||||
from app.database import Base, SessionLocal, engine, get_db
|
||||
from app.main import app
|
||||
from tools import dbmeter
|
||||
from tools.fakeprose import prose
|
||||
|
||||
# A stand-in for the real thing: the assembled prompt, which is what makes the
|
||||
# column enormous, plus the small world_state slice the UI actually needs.
|
||||
#
|
||||
# Varied text, not `"x" * 20_000`. The column is stored compressed now
|
||||
# (migration 43), and a repeated character compresses about a thousandfold —
|
||||
# which would make the byte ceilings below pass against a fixture that costs
|
||||
# nothing, testing nothing. Prose-shaped filler compresses like the prompts
|
||||
# this stands in for.
|
||||
_SNAPSHOT_RNG = random.Random(20_260_817)
|
||||
BIG_SNAPSHOT = {
|
||||
"system": "x" * 20_000,
|
||||
"story": "y" * 40_000,
|
||||
"system": prose(_SNAPSHOT_RNG, 20_000),
|
||||
"story": prose(_SNAPSHOT_RNG, 40_000),
|
||||
"world_state": {
|
||||
"delta": {"player.hp": -15},
|
||||
"report": {"applied": [{"path": "player.hp", "old": 100, "new": 85}]},
|
||||
@@ -204,16 +213,37 @@ def test_snapshot_is_still_reachable_on_demand(client):
|
||||
action_id = r.json()["actions"][0]["id"]
|
||||
r = client.get(f"/api/adventures/{client.adv_id}/actions/{action_id}/context")
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["system"] == "x" * 20_000
|
||||
# Round-tripped through zlib and back to a dict, byte for byte.
|
||||
assert r.json()["system"] == BIG_SNAPSHOT["system"]
|
||||
assert r.json()["story"] == BIG_SNAPSHOT["story"]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ backfill
|
||||
|
||||
def as_json_snapshot_column(db) -> None:
|
||||
"""Put actions.context_snapshot back as JSON, the way it was before 43.
|
||||
|
||||
Migration 36 lifts world_delta out of the snapshot with SQL JSON
|
||||
functions, so it can only run while the column still *is* JSON. In a real
|
||||
upgrade it always is — 36 runs seven migrations before 43 compresses the
|
||||
column into a BLOB — but `create_all` builds today's schema, so a test
|
||||
calling that backfill has to rebuild the schema it was written against.
|
||||
"""
|
||||
db.execute(text("ALTER TABLE actions DROP COLUMN context_snapshot"))
|
||||
db.execute(text("ALTER TABLE actions ADD COLUMN context_snapshot JSON"))
|
||||
db.execute(
|
||||
text("UPDATE actions SET context_snapshot = :snapshot"),
|
||||
{"snapshot": json.dumps(BIG_SNAPSHOT)},
|
||||
)
|
||||
db.commit()
|
||||
|
||||
|
||||
def test_backfill_populates_world_delta_from_existing_snapshots(client):
|
||||
"""Migration 36 lifts the slice out server-side, without reading the
|
||||
snapshots into Python."""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
as_json_snapshot_column(db)
|
||||
db.execute(text("UPDATE actions SET world_delta = NULL"))
|
||||
db.commit()
|
||||
assert db.query(models.Action).filter(models.Action.world_delta.isnot(None)).count() == 0
|
||||
@@ -419,9 +449,14 @@ def test_the_ceiling_discriminates(client, meter):
|
||||
"""A ceiling is only worth having if the thing it excludes would breach it.
|
||||
|
||||
This is the regression the byte tests exist to catch, performed on purpose:
|
||||
undefer the snapshot and the same twelve rows cost two orders of magnitude
|
||||
more. If this ever stops exceeding the budget, the fixture has gone too
|
||||
small for the tests above to mean anything.
|
||||
undefer the snapshot and the same twelve rows cost several times the whole
|
||||
budget. If this ever stops exceeding it, the fixture has gone too small for
|
||||
the tests above to mean anything.
|
||||
|
||||
The margin used to be a hundredfold and is now about six. That is not the
|
||||
guard weakening — it is migration 43 compressing the column, and the
|
||||
fixture text being prose-shaped so it compresses like a real prompt rather
|
||||
than like a repeated character.
|
||||
"""
|
||||
budget = ACTIONS_IN_FIXTURE * PAGE_LOAD_BYTES_PER_ACTION
|
||||
db = SessionLocal()
|
||||
@@ -436,8 +471,8 @@ def test_the_ceiling_discriminates(client, meter):
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
assert fetched(meter) > budget * 10, (
|
||||
assert fetched(meter) > budget * 3, (
|
||||
"undeferring the snapshot cost only "
|
||||
f"{fetched(meter):,} B — the fixture is too small for the byte "
|
||||
"ceilings above to catch anything"
|
||||
f"{fetched(meter):,} B against a {budget:,} B budget — the fixture is "
|
||||
"too small for the byte ceilings above to catch anything"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,260 @@
|
||||
"""context_snapshot, stored compressed.
|
||||
|
||||
The column is 89% of the database and the free tier allows 512 MB. Reads were
|
||||
solved by deferring it; this is about the storage ceiling. Postgres already
|
||||
TOASTs it and only gets 1.7x, because pglz is tuned for fast decompression of
|
||||
data a query might filter on — and nothing ever filters on an assembled
|
||||
prompt.
|
||||
|
||||
Three things have to hold, and only the first is obvious:
|
||||
|
||||
* what goes in comes back out, exactly, including a snapshot written before
|
||||
the conversion and one that is NULL;
|
||||
* the model still hands callers a dict, so no call site changes;
|
||||
* migration 44 drops the original column, so the backfill is the one
|
||||
destructive step in this file — it must convert every row or abort.
|
||||
|
||||
python -m pytest tests/test_snapshot_compression.py -v
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import tempfile
|
||||
import zlib
|
||||
|
||||
_tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
|
||||
_tmp.close()
|
||||
os.environ["AIDND_DB_PATH"] = _tmp.name
|
||||
os.environ.pop("AIDND_DATABASE_URL", None)
|
||||
os.environ.pop("DATABASE_URL", None)
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from app import compression, migrations, models
|
||||
from app.database import Base, SessionLocal, engine
|
||||
from tools.fakeprose import prose
|
||||
|
||||
|
||||
def snapshot(seed: int, nbytes: int = 60_000) -> dict:
|
||||
rng = random.Random(seed)
|
||||
return {
|
||||
"system": prose(rng, nbytes // 3),
|
||||
"story": prose(rng, nbytes - nbytes // 3),
|
||||
"world_state": {"delta": {"player.hp": -3}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def db():
|
||||
Base.metadata.create_all(bind=engine)
|
||||
session = SessionLocal()
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
session.close()
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def adventure(db):
|
||||
user = models.User(is_guest=False, email="snap@example.com")
|
||||
db.add(user)
|
||||
db.flush()
|
||||
adv = models.Adventure(user_id=user.id, title="Cave", script_state={})
|
||||
db.add(adv)
|
||||
db.commit()
|
||||
return adv
|
||||
|
||||
|
||||
# ------------------------------------------------------------------- packing
|
||||
|
||||
def test_pack_round_trips_exactly():
|
||||
value = snapshot(1)
|
||||
assert compression.unpack(compression.pack(value)) == value
|
||||
|
||||
|
||||
def test_pack_handles_the_awkward_values():
|
||||
for value in ({}, {"a": None}, {"nested": {"deep": [1, 2, {"x": "é"}]}}):
|
||||
assert compression.unpack(compression.pack(value)) == value
|
||||
|
||||
|
||||
def test_pack_actually_shrinks_prose():
|
||||
"""The whole justification. If this ratio collapses, the migration is
|
||||
spending CPU for nothing."""
|
||||
value = snapshot(2, 200_000)
|
||||
raw = json.dumps(value, separators=(",", ":")).encode()
|
||||
packed = compression.pack(value)
|
||||
assert len(packed) < len(raw) / 3, (
|
||||
f"{len(raw):,} B compressed to {len(packed):,} B — under 3x"
|
||||
)
|
||||
|
||||
|
||||
def test_unpack_rejects_nothing_it_wrote():
|
||||
packed = compression.pack({"a": "b"})
|
||||
assert compression.unpack(bytearray(packed)) == {"a": "b"}
|
||||
assert compression.unpack(memoryview(bytes(packed))) == {"a": "b"}
|
||||
|
||||
|
||||
# --------------------------------------------------------------- the column
|
||||
|
||||
def test_the_column_stores_bytes_and_returns_a_dict(db, adventure):
|
||||
value = snapshot(3)
|
||||
action = models.Action(
|
||||
adventure_id=adventure.id, index=0, type="ai", text="t",
|
||||
context_snapshot=value,
|
||||
)
|
||||
db.add(action)
|
||||
db.commit()
|
||||
db.expire_all()
|
||||
|
||||
assert action.context_snapshot == value
|
||||
|
||||
stored = db.execute(
|
||||
text("SELECT context_snapshot FROM actions WHERE id = :id"),
|
||||
{"id": action.id},
|
||||
).scalar()
|
||||
assert isinstance(stored, (bytes, bytearray, memoryview))
|
||||
assert json.loads(zlib.decompress(bytes(stored))) == value
|
||||
|
||||
|
||||
def test_the_column_is_smaller_than_the_json_it_holds(db, adventure):
|
||||
value = snapshot(4, 200_000)
|
||||
action = models.Action(
|
||||
adventure_id=adventure.id, index=0, type="ai", text="t",
|
||||
context_snapshot=value,
|
||||
)
|
||||
db.add(action)
|
||||
db.commit()
|
||||
|
||||
stored = db.execute(
|
||||
text("SELECT length(context_snapshot) FROM actions WHERE id = :id"),
|
||||
{"id": action.id},
|
||||
).scalar()
|
||||
assert stored < len(json.dumps(value)) / 3
|
||||
|
||||
|
||||
def test_null_stays_null(db, adventure):
|
||||
action = models.Action(
|
||||
adventure_id=adventure.id, index=0, type="do", text="t",
|
||||
context_snapshot=None,
|
||||
)
|
||||
db.add(action)
|
||||
db.commit()
|
||||
db.expire_all()
|
||||
assert action.context_snapshot is None
|
||||
|
||||
|
||||
def test_an_unreadable_snapshot_reads_as_none_rather_than_raising(db, adventure):
|
||||
"""One corrupt row must not 500 the turn that happens to load it. The
|
||||
snapshot is a debugging view; the story is the thing that matters."""
|
||||
action = models.Action(
|
||||
adventure_id=adventure.id, index=0, type="ai", text="t",
|
||||
context_snapshot={"a": "b"},
|
||||
)
|
||||
db.add(action)
|
||||
db.commit()
|
||||
db.execute(
|
||||
text("UPDATE actions SET context_snapshot = :junk WHERE id = :id"),
|
||||
{"junk": b"not zlib at all", "id": action.id},
|
||||
)
|
||||
db.commit()
|
||||
db.expire_all()
|
||||
|
||||
assert db.get(models.Action, action.id).context_snapshot is None
|
||||
|
||||
|
||||
# ------------------------------------------------------- the upgrade in full
|
||||
|
||||
def as_json_column(db, rows: dict[int, dict]) -> None:
|
||||
"""The pre-43 schema: context_snapshot as a JSON column, populated."""
|
||||
db.execute(text("ALTER TABLE actions DROP COLUMN context_snapshot"))
|
||||
db.execute(text("ALTER TABLE actions ADD COLUMN context_snapshot JSON"))
|
||||
for action_id, value in rows.items():
|
||||
db.execute(
|
||||
text("UPDATE actions SET context_snapshot = :v WHERE id = :id"),
|
||||
{"v": json.dumps(value), "id": action_id},
|
||||
)
|
||||
db.commit()
|
||||
|
||||
|
||||
def seed_pre_43(db, adventure, count: int = 4) -> dict[int, dict]:
|
||||
ids = []
|
||||
for i in range(count):
|
||||
action = models.Action(
|
||||
adventure_id=adventure.id, index=i, type="ai", text=f"t{i}"
|
||||
)
|
||||
db.add(action)
|
||||
db.flush()
|
||||
ids.append(action.id)
|
||||
# One action with no snapshot at all, which must survive as NULL.
|
||||
plain = models.Action(
|
||||
adventure_id=adventure.id, index=count, type="do", text="look"
|
||||
)
|
||||
db.add(plain)
|
||||
db.commit()
|
||||
|
||||
expected = {action_id: snapshot(action_id) for action_id in ids}
|
||||
as_json_column(db, expected)
|
||||
db.execute(text(f"PRAGMA user_version = {migrations.SNAPSHOT_COMPRESS_VERSION - 1}"))
|
||||
db.commit()
|
||||
return expected
|
||||
|
||||
|
||||
def test_bootstrap_converts_every_snapshot(db, adventure):
|
||||
expected = seed_pre_43(db, adventure)
|
||||
db.close()
|
||||
|
||||
migrations.bootstrap(engine)
|
||||
|
||||
check = SessionLocal()
|
||||
try:
|
||||
for action_id, value in expected.items():
|
||||
assert check.get(models.Action, action_id).context_snapshot == value
|
||||
assert check.execute(
|
||||
text("SELECT count(*) FROM actions WHERE context_snapshot IS NULL")
|
||||
).scalar() == 1
|
||||
finally:
|
||||
check.close()
|
||||
|
||||
|
||||
def test_bootstrap_leaves_the_column_named_context_snapshot(db, adventure):
|
||||
seed_pre_43(db, adventure)
|
||||
db.close()
|
||||
|
||||
migrations.bootstrap(engine)
|
||||
|
||||
with engine.begin() as conn:
|
||||
columns = {row[1] for row in conn.execute(text("PRAGMA table_info(actions)"))}
|
||||
assert "context_snapshot" in columns
|
||||
assert "context_snapshot_z" not in columns, "the swap left the scratch column behind"
|
||||
|
||||
|
||||
def test_the_backfill_aborts_rather_than_dropping_unconvertible_data(
|
||||
db, adventure, monkeypatch
|
||||
):
|
||||
"""Migration 44 destroys the original. If anything cannot be converted the
|
||||
whole run has to roll back with the column still there — the alternative is
|
||||
losing somebody's prompts to a bug in this file."""
|
||||
seed_pre_43(db, adventure)
|
||||
db.close()
|
||||
|
||||
def broken_unpack(_blob):
|
||||
return {"not": "what went in"}
|
||||
|
||||
monkeypatch.setattr(compression, "unpack", broken_unpack)
|
||||
|
||||
with pytest.raises(RuntimeError, match="round trip"):
|
||||
migrations.bootstrap(engine)
|
||||
|
||||
with engine.begin() as conn:
|
||||
columns = {row[1] for row in conn.execute(text("PRAGMA table_info(actions)"))}
|
||||
version = conn.execute(text("PRAGMA user_version")).scalar()
|
||||
surviving = conn.execute(
|
||||
text("SELECT count(*) FROM actions WHERE context_snapshot IS NOT NULL")
|
||||
).scalar()
|
||||
|
||||
assert "context_snapshot" in columns
|
||||
assert version < migrations.SNAPSHOT_COMPRESS_VERSION, \
|
||||
"the version advanced past a backfill that failed"
|
||||
assert surviving == 4, "the prompts did not survive the rollback"
|
||||
Reference in New Issue
Block a user