Skip to content

MAINT: Batch the seed dedupe lookup when adding seeds to memory - #2523

Open
varunj-msft wants to merge 1 commit into
microsoft:mainfrom
varunj-msft:varunj-msft/v1.1.0-Release-Seed-Dedupe-Batching
Open

MAINT: Batch the seed dedupe lookup when adding seeds to memory#2523
varunj-msft wants to merge 1 commit into
microsoft:mainfrom
varunj-msft:varunj-msft/v1.1.0-Release-Seed-Dedupe-Batching

Conversation

@varunj-msft

Copy link
Copy Markdown
Contributor

Description

add_seeds_to_memory_async ran one SELECT per seed to decide whether that seed was already stored. value_sha256 is not indexed, so each of those scans the table, and the cost grows with the number of rows already present — which makes the second load of a dataset far more expensive than the first.

Loading the default datasets adds roughly 158,000 seeds, and this lookup dominated the entire load. Measured on main against a real database:

  • First load, empty table growing to 158,580 rows: about 210 seconds.
  • Second load, table already holding 158,580 rows: over 1140 seconds, and still running when I stopped it.

LoadDefaultDatasets.initialize_async has no already-loaded short-circuit, so it re-resolves the whole selection every time it runs. A session-scoped backend serving many scenarios therefore pays that second, much worse, cost repeatedly.

This prepares the seeds first, then resolves which hashes are already stored using chunked lookups rather than one query per seed, so a full dataset resolves in a handful of queries instead of roughly 158,000. The chunk size of 500 is deliberately well under SQL Server's 2100-parameter ceiling.

Existing behaviour is preserved deliberately, not incidentally:

  • The stored keys are snapshotted once, before any comparison, so duplicates within a single call are all still inserted. Only what storage held when the call started counts as already present. This matches what the per-seed loop did, since nothing it inserted was visible to its own subsequent lookups either.
  • A seed with no dataset name still matches its hash in any dataset, mirroring the unfiltered lookup it previously performed.

There is no schema change and no migration here. Adding a bounded type and an index on value_sha256 would be the more complete fix, but that needs an Alembic revision, and alembic check runs on every memory init, so it is intentionally left as a follow-up rather than landing during a release window. This change is purely query shape and requires no database changes at all.

Part of the v1.1.0 release wave with #2510, #2511 and #2512.

Tests and Documentation

Four new tests in tests/unit/memory/memory_interface/test_interface_seed_prompts.py:

  • test_add_seed_prompts_duplicates_within_one_call_are_all_stored pins the snapshot semantics, which is the subtle way a change here could silently alter behaviour.
  • test_add_seed_prompts_without_dataset_name_matches_any_dataset covers the unfiltered lookup path.
  • test_add_seed_prompts_dedupes_across_chunk_boundaries uses 1025 seeds, two full chunks plus a partial, then re-adds them with one new seed to confirm only the new one is stored.
  • test_add_seed_prompts_queries_are_batched is a regression guard: it asserts get_seeds is called exactly once for 50 seeds, so a revert to per-seed lookups fails loudly instead of quietly getting slow again.

The existing duplicate-handling tests, same dataset and different datasets, are unchanged and still pass, which is the main evidence that this is a pure performance change.

I also diffed old versus new results across seven edge cases, including 2500 hashes in a single call to confirm the chunking respects the parameter ceiling. The retained and skipped sets were identical in every case.

Ran pytest tests/unit/memory/memory_interface/test_interface_seed_prompts.py: 70 passed.

No documentation changes: the public API and its behaviour are unchanged. The add_seeds_to_memory_async docstring now states the within-call duplicate rule explicitly, which was previously only implied. JupyText was not run and is not applicable: no notebooks or code samples are affected.

@hannahwestra25 hannahwestra25 self-assigned this Sep 1, 2026
Comment thread pyrit/memory/memory_interface.py Outdated
Comment thread pyrit/memory/memory_interface.py
@varunj-msft
varunj-msft force-pushed the varunj-msft/v1.1.0-Release-Seed-Dedupe-Batching branch 2 times, most recently from 7da86bb to 5940b16 Compare September 1, 2026 19:30
Comment thread pyrit/memory/memory_interface.py Outdated
Comment thread pyrit/memory/memory_interface.py Outdated
add_seeds_to_memory_async ran one SELECT per seed to decide whether that seed
was already stored. value_sha256 is not indexed, so each of those scans the
table, and the cost grows with the number of rows already present. Loading the
default datasets adds roughly 158,000 seeds, where this dominated the whole
load, and it was worse still on every later load because the table was already
full.

Prepare the seeds first, then resolve which hashes are already stored in chunked
lookups instead of per seed. get_seeds issues its statement directly rather than
through the batching helpers, so the lookup has to bound the IN clause itself. It
chunks on _MAX_BIND_VARS, the per-statement bind ceiling each backend already
tunes for itself, so SQL Server gets the 2000 AzureSQLMemory sets rather than a
separate hardcoded limit.

The lookups are grouped by dataset name so the name is still compared by the
database. Equality there is a property of the column's collation: Azure SQL's
default is case-insensitive and also ignores trailing blanks, so an existing
("Dataset") row matched an incoming ("dataset") one before this change.
Comparing the names in Python would impose one fixed rule on every backend and
insert a duplicate wherever the stored spelling differed, which no test would
catch because CI runs on SQLite, whose default collation is case-sensitive.
Since the database decides the match, the pair is keyed by the requested name
rather than the stored one. An empty name filters nothing, exactly like None, so
the two are normalized together to keep the lookup and the comparison agreeing.

The stored keys are snapshotted once, before any comparison, which preserves the
existing behaviour that duplicates within a single call are all inserted; only
what storage held when the call started counts as already present. Seeds without
a dataset name still match their hash in any dataset, matching the unfiltered
lookup they previously performed.

The dataset name takes one of those binds, so the hashes get the ceiling minus one
rather than the whole budget. Binding a full chunk plus the name would have gone
one over the limit each backend declares.
@varunj-msft
varunj-msft force-pushed the varunj-msft/v1.1.0-Release-Seed-Dedupe-Batching branch from 5940b16 to 345829a Compare September 1, 2026 23:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants