[ISSUE-844] fix(ai): invalidate an embedding vector when its text changes - #858
[ISSUE-844] fix(ai): invalidate an embedding vector when its text changes#858E2ern1ty wants to merge 2 commits into
Conversation
…nges EmbeddingIndexStore keyed each persisted vector by the entity key, which is a prefix, the id and the label, with nothing derived from the value that was embedded. initStore treated the presence of that key as "this entity is indexed", so an entity that kept its id and changed its value silently kept the vector of the value it no longer had. Retrieval for the old value still reached the entity while its current value was unreachable, and since the subgraph handed on is verbalized from the current graph, a query about the old value came back with a context showing the new one. A record now carries a fingerprint of the text it was produced from, a new contentHash field on EmbeddingResult holding SHA-256 over the entity's chunk list. On rebuild a record is loaded only when that fingerprint matches what the entity would be embedded from now; a mismatch counts as not indexed, so the entity is embedded again. Both paths derive the chunk list through one method, so the text hashed on the way in is the text hashed on the way out, which takes verbalize(GraphEntity) to be a function of the entity alone. The number of records dropped is logged, since the complaint in the issue is not only that the vector was stale but that nothing said so. Existing index files. A record without a fingerprint cannot be shown to match the current value, and trusting it would keep exactly the staleness this check is for, so it is not loaded and its entity is embedded once. That is one of the two treatments the issue put up for decision, taken because the other one, valid until the entity is next touched, never comes due: nothing touches such a record. The file is still only appended to, and the second half of the issue is scoped as the issue itself suggests. A superseded record stays in the file and is rejected on each later read, so the file grows by one record set per distinct value an entity has held, and a value that comes back is served from the record already there without a request. The records of deleted entities also stay. Pruning that history wants a rewrite rather than an append, which the issue puts down as separate work, better done as compaction. The harm named there does not wait for it: an id taken over by another entity no longer adopts the vector left behind, because that vector's fingerprint does not match the new text. EmbeddingIndexInvalidationTest drives the store against a local embeddings endpoint answering with one hot vectors, one dimension per distinct text, so a stored vector says which text produced it. Six cases: a changed value on a vertex and on an edge, recall through GraphMemoryServer going the right way on both the old and the current value, an id taken over by another entity, a value that comes back, and a record without a fingerprint. Each fails against the behaviour it pins. The unchanged run asserts the file is byte identical, which is what would catch a fingerprint computed differently on the two paths. GraphMemoryTest no longer loads an embedding store from a checked in index file, and that file is removed. It could not be carried over: its records line up with the text the current code would embed for only 14 of its 163 entities, the rest disagreeing on chunk count, so there was nothing honest to stamp them with and no way to embed them again offline. It also could not have affected a single assertion, since every query there embeds to double[0] through MockChatRobot and EmbeddingVector.match returns 0.0 on a length mismatch, below the 0.5 threshold. Its assertions come from the keyword path and are unchanged; the store's load path is now covered deliberately by the new test. Not addressed here. The fingerprint covers the text only, so changing the embedding model or its dimension leaves every record loaded, which is the same class of staleness; which properties of a model identify it is a separate decision. And a record set short of a chunk still matches on every record it has, so an entity can be held with part of its vectors, as it could before, since a fingerprint covers all of an entity's chunks at once and carries no ordinal.
f8511d1 to
61bd968
Compare
| unversionedRecords++; | ||
| } else if (embedding.contentHash.equals( | ||
| currentFingerprint(key, entity, key2Fingerprint))) { | ||
| this.indexStoreMap.computeIfAbsent(entity, k -> new ArrayList<>()).add(embedding); |
There was a problem hiding this comment.
I've observed a phenomenon: the code uses a Map<GraphEntity, List<EmbeddingResult>>(indexStoreMap) and populates it with entity instances in the initStore; however, subsequent locations can access the index using other GraphAccessor/GraphEntity instances (e.g., the GraphEntity instance generated when getEntityIndex is called externally might not be the same object reference). If the GraphEntity lacks appropriate equals/hashCode or object instance adjacency, it can lead to not finding loaded records (a logical error); furthermore, using a GraphEntity object as a long-term key maintains a strong reference to the graph, causing memory leaks.
Relevant code:
// Snippet from EmbeddingIndexStore.initStore (already exists in the PR)
GraphEntity entity = key2EntityMap.get(key);
if (entity != null) {
this.indexStoreMap.computeIfAbsent(entity, k -> new ArrayList<>()).add(embedding);
}
Suggested solution (prioritize stable string keys instead of objects): Change the indexStoreMap to a Map<String, List> with the entity key (ModelUtils.getGraphEntityKey(entity)) as the key, and replace the key with this string in all reads. This avoids memory issues caused by referencing entity objects and ensures consistency across Accessors/instances.
There was a problem hiding this comment.
Done in 923347f, keyed by ModelUtils.getGraphEntityKey as you suggest, though I want to correct the two reasons given, because they point at the wrong risk and I would rather the record be accurate.
Lookups across instances already worked. GraphVertex.equals/hashCode delegate to Vertex, which is Objects.equals(id) && Objects.equals(label), and GraphEdge to Edge, which is srcId, dstId and label. So an instance handed out by a different GraphAccessor for the same entity is equal and hashes alike, and getEntityIndex found the records. I could not construct a case on master where it did not.
What a key retained was not the graph. GraphVertex holds a Vertex, and Vertex holds id, label and values — no reference to MemoryGraph, GraphAccessor or anything above it. So the retention was of the Vertex objects and their value strings, which is real but bounded by the entity count and small beside the vectors in the same map, at 1024 doubles each. Not a leak in the sense of holding a graph alive.
The reason I think the change is right is narrower, and closer to what this PR is about. With the entity as the key, what counts as the same indexed entity is decided by GraphVertex.equals, which answers on the id and the label. This PR exists because an entity's identity did not cover its value. Extending Vertex.equals to cover the values is a reasonable thing for someone to want next, and the moment they did, getEntityIndex would stop finding records and the scan in initStore would stop skipping what is already indexed — silently, both of them. Meanwhile the file has always been keyed by the entity key. Two notions of identity for one thing, one of them defined in a class this PR invites people to change. Keying the map by the same string as the file collapses them into one and takes the index out of reach of that change.
The call sites moved with it: the scan loop and the in-batch dedup set now hold keys, indexBatch writes under the key it already computes for res.input, and getEntityIndex looks up getGraphEntityKey(entity). That last one builds a small string per lookup where before it hashed two fields, which on the query path is one allocation per vertex per query, against a 1024-dimension cosine for each of them.
There is a test for the scenario you describe: it indexes through one graph and reads through an entity object from another, asserting first that the two objects really are different. It holds on master as well — it is there to keep it holding.
| String key = embedding.input; | ||
| GraphEntity entity = key2EntityMap.get(key); | ||
| if (entity != null) { | ||
| this.indexStoreMap.computeIfAbsent(entity, k -> new ArrayList<>()).add(embedding); |
There was a problem hiding this comment.
If the indexStoreMap is read and written by different threads during runtime (e.g., index initialization on the startup thread, querying on the worker thread), concurrent read/write safety must be guaranteed. Replace the map with a ConcurrentHashMap.
There was a problem hiding this comment.
Addressed in 923347f, but not with a ConcurrentHashMap, because I do not think that fixes the case you describe.
Take it as stated: initialization on the startup thread, a query on a worker thread. initStore assigned the map to the field up front and then mutated it for the whole of its run — the file read, then batch after batch of embedding. A worker reading in that window sees whatever has been filled in so far. With a ConcurrentHashMap it still does: every individual get is safe, no exception is thrown, and the answer is a partial index. For a recall path that is worse than an error, because it comes back as a plausible result rather than a failure. The same goes for a second initStore on a live store, which under either map empties the index for readers before filling it again.
So the map is built aside and assigned to the field in one go, at the end of initStore. The field is volatile, and starts as Collections.emptyMap(). A reader now sees either the index that was there before or the finished one, and a reader that arrives before any build has happened gets an empty list rather than the NullPointerException it would have got on master. getEntityIndex reads the field once per call, so it works against one map throughout.
Two notes on the trade-off, since it is not free:
- During a rebuild both maps are held, so peak memory is roughly doubled for that window. I think consistency is worth it here, and there is no caller in tree that rebuilds a live store — nothing in
geaflow-ai/src/mainconstructs this store at all today. - This does not make
initStoresafe to call concurrently with itself. Two builds at once still race on the index file, which is appended to. That was true before and I have not tried to fix it here; if you would like it enforced, the smallest thing would be to reject a concurrent call outright rather than to make it work.
There is a test: it holds the embeddings endpoint open with a latch, runs a second initStore on the same store on another thread, and asserts that a reader in that window gets the vector from before the build. It fails if the field is assigned before the build instead of after — with the assignment moved back, the reader gets zero vectors, which is the partial index a ConcurrentHashMap would also have shown.
…ublish it once Review of apache#858 raised two things about indexStoreMap. Both are about the map rather than the fingerprint, and both are addressed, though not quite as suggested. Keyed by a string, not by the entity. The map is now Map<String, List<EmbeddingResult>> under ModelUtils.getGraphEntityKey, the same string the index file is keyed by, so one notion of which entity a vector belongs to serves both memory and disk. Lookups across instances already worked, since GraphVertex.equals answers on the id and the label and GraphEdge on the two ids and the label, and what a key retained was the Vertex and its values rather than the graph. The reason to change it is narrower and more to the point of this fix: with the entity as the key, what counts as the same indexed entity is decided by GraphVertex.equals, and this is a fix about an entity's identity not covering its value. Someone extending that equals to cover the values, which is a reasonable thing to want, would silently stop getEntityIndex finding records and stop the scan skipping what is indexed. Keying by the entity key takes the index out of the reach of that change. Published once, rather than made concurrent. initStore built into the field and mutated it throughout, so a reader on another thread saw an index half way through being built, and a second initStore emptied it for readers before filling it again. A ConcurrentHashMap makes each operation safe but leaves that reader looking at a partial index, which is the wrong answer rather than a crash. So initStore now builds into a local map and assigns it to the field in one go, the field is volatile and starts empty, and a reader is shown either the index that was there before or the finished one, and no longer has to have waited for a first build to avoid failing. Two cases added. One reads through an entity object from a different graph than the one that was indexed, which is the scenario raised; it holds on master too and is there to keep it holding. The other holds the endpoint open, reads while a second build is in flight, and requires the vector from before the build; it fails if the field is assigned before the build rather than after.
2 participants
What changes were proposed in this pull request?
Fixes #844.
EmbeddingIndexStorekeyed each persisted vector by the entity key, a prefix plus the id and the label, with nothing derived from the value that was embedded.initStoreread the presence of that key as "this entity is indexed", so an entity that kept its id and changed its value silently kept the vector of the value it no longer had. Retrieval for the old value still reached the entity while its current value was unreachable, and because the subgraph handed on is verbalized from the current graph, a query about the old value came back with a context showing the new one.A record now says which text produced it.
EmbeddingResultgains acontentHashfield, SHA-256 over the entity's chunk list. On rebuild a record is loaded only when that fingerprint matches what the entity would be embedded from now; a mismatch counts as not indexed, so the entity is embedded again. Both paths derive the chunk list through one method, so the text hashed on the way in is the text hashed on the way out. The number of records dropped is logged, since the complaint in the issue is not only that the vector was stale but that nothing said so.Existing index files are treated as the first of the two options the issue put up for decision. A record without a fingerprint cannot be shown to match the current value, and trusting it would keep exactly the staleness the check is for, so it is not loaded and its entity is embedded once. The other option, valid until the entity is next touched, never comes due, because nothing touches such a record.
The production change is +149 lines across three files, of which about 45 come from the review response below rather than from the fix itself.
Measured against a real embedding model
The same harness was run on
master(3f73eb5) and on this branch, both driving the store at BAAI/bge-m3, 1024 dimensions, through a local proxy that counts requests and records what is sent. The two values used are unrelated:cosine(embed(old), embed(new)) = 0.2923, so a hit on one is not a hit on both.A. The value of an entity changes, id and label unchanged.
The master column reproduces the issue exactly, including the part a fact-checking layer cannot see from the outside: the entity is retrieved because of the old value and the reader is then shown the new one. The issue reported 1.0000 and 0.2357 for those two cosines against its own pair of values; here the same measurement gives 1.0000 and 0.2923.
B. An entity is removed and another takes its id over, which is the harm the issue names for records that are never pruned. Three runs: indexed, gone, id reused with a different value.
So on master the new entity silently adopts the vector left behind. It no longer does, and that does not wait for compaction.
C. A value the entity has held before comes back. On this branch the third run issues 0 requests and the file stays at 2 lines: the record already there matches again. This is what bounds the growth of an append-only file to one record set per distinct value.
D. A record with no fingerprint, holding a vector that happens to be correct. On master it is kept and nothing is sent. On this branch the first run sends 1 request and the second sends 0 — the cost of the format change is paid once, per entity, and never again.
Nothing about the credentials used is in the diff; the harness read them from the environment and is not part of this PR.
Review response, second commit
Two things were raised about
indexStoreMap, both about the map rather than the fingerprint. Addressed in 923347f, though not quite as suggested; the reasoning is in the review threads.ModelUtils.getGraphEntityKeyinstead of by the entity. Lookups across instances already worked, sinceGraphVertex.equalsanswers on the id and the label, and what a key retained was theVertexand its values rather than the graph. The reason to change it is narrower: with the entity as the key, what counts as the same indexed entity is decided byGraphVertex.equals, and this is a fix about an entity's identity not covering its value. Someone extending thatequalsto cover the values would silently stopgetEntityIndexfinding records and stop the scan skipping what is indexed. The file has always been keyed by the entity key; now the map is too.initStorebuilt into the field and mutated it throughout, so a reader on another thread saw an index half way through being built. AConcurrentHashMapmakes each operation safe but leaves that reader looking at a partial index, which on a recall path is worse than an error. SoinitStorebuilds into a local map and assigns it in one go, the field isvolatileand starts empty, and a reader sees either the index that was there before or the finished one — and no longer needs a build to have happened to avoid an NPE. This does not make concurrentinitStorecalls safe against each other; they still race on the file, as before.What is deliberately left out
The file is still only appended to, which is where the issue puts the line: "Pruning orphan records wants a rewrite of the file rather than an append, which is a separate piece of work from invalidation and may be better done as compaction."
So a superseded record stays in the file and is rejected on each later read, and the records of deleted entities stay too. Measurements C and B above are why that is liveable: growth is bounded by the number of distinct values an entity has held rather than by the number of changes, and the dangerous half of the pruning problem is already closed by the fingerprint. What remains for compaction is only that the file size does not track the graph.
I had a rewrite-on-load in the first revision of this PR and took it out. It cost a public
upgradeIndexFile, a temp-file-and-fsync write path, and the bookkeeping to keep a graph that scanned short from taking the whole file with it — about 220 lines to prune history the issue had already scoped out.One thing to look at: a test fixture is removed
GraphMemoryTestno longer loads an embedding store fromsrc/test/resources/index/LDBCEmbeddingIndexStore, and that 4.6MB file is deleted. Two findings led there, both checkable before the deletion:double[0]throughMockChatRobot, andEmbeddingVector.matchreturns0.0on a length mismatch, below the0.5threshold. Everything asserted there comes from the keyword path, and those 22 assertions pass unchanged. Runtime for that test drops from 27.6s to 0.6s, the 27s having been model retries against a null URL.If you would rather keep the file, I can drop only the wiring, though it would then be unreferenced. The
**/resources/index/**rat exclusion in the root pom is now unused; I left it alone as it is a generic pattern, and can remove it on request.Two known limits, neither introduced here
indexBatch. This is unchanged from before: a partial set was skipped as indexed then too.The verify path also takes
verbalize(GraphEntity)to be a function of the entity alone, noted in the javadoc where it is relied on. The one implementation in tree is.How was this PR tested?
Verified against a real hosted embedding model as set out above, though not in a production deployment, hence the second box.
EmbeddingIndexInvalidationTestcovers the same ground offline, driving the store against a local/v1/embeddingsendpoint that answers with one hot vectors, one dimension per distinct text, so a stored vector says which text produced it: cosine 1 against its own text and 0 against any other. No key and no network needed. Eight cases:GraphMemoryServerEach case fails against the behaviour it pins, checked by disabling the fingerprint comparison and, for the last one, by assigning the index before the build instead of after. The byte-identical assertion on the unchanged run is what would catch a fingerprint computed differently on the two paths.
mvn -pl geaflow-ai clean test -Pjdk8: 14 tests, 0 failures, checkstyle and RAT clean. The one case with threads in it was run repeatedly to check it settles. Nothing ingeaflow-ai/src/mainconstructs this store today, so the added per-entity hashing on load has no production caller in tree.