Auditing Whether a Re-Embedding Migration Actually Finished

#ai #machinelearning #tutorial #llm

Re-embedding jobs do not usually fail. They stop — a worker is rescheduled, a partition is skipped, a retry loop gives up on a page — and what is left behind is an index that answers every query and is missing eleven thousand documents.

Tag every vector at write time

None of the checks below are possible unless each stored vector carries enough metadata to say where it came from. Five fields, written by the embedding job on every record:

embed_model — the exact model identifier sent to the provider, not a config alias.
embed_dim — the output width actually requested.
embed_run_id — an identifier for this migration run. This is what lets you distinguish “written by the backfill” from “written by live ingestion while the backfill was running”, and those two populations behave differently.
source_version — the version, revision or content hash of the source record that was embedded. Not a timestamp of the embedding: the version of the input.
embedded_at — when the vector was produced.

These are cheap. On a corpus of ten million chunks they are tens of megabytes of metadata against tens of gigabytes of vectors, and they are the difference between an audit and an opinion. If you are only adding them now, add them as part of the re-embed itself — the records being rewritten are exactly the ones that need them.

Why row counts prove nothing

The check everybody writes first is: count the source rows, count the destination vectors, assert equal. It passes on jobs that are badly broken, for a reason that is structural rather than unlucky.

A batch job that retries is a job that can write the same record twice. A worker that upserts a page of two hundred records, dies before committing its cursor, and is restarted will rewrite those two hundred. If the write is a true upsert keyed on identifier, that is harmless — but if the destination assigns its own identifiers, or the job derives a new identifier per run, the duplicates are separate records. Now the destination contains duplicates and is missing the page the job skipped when it later crashed elsewhere, and the two errors cancel in the total. Equal counts. Broken index.

Counts also cannot see the failure that matters most, which is scope: a job that processed one namespace of four and terminated cleanly. Counted globally against a source query that was also scoped to the same namespace, that job passes. Any count check must therefore be grouped by every partitioning dimension you have — namespace, tenant, shard, document type — and asserted per group, with the group list taken from the source rather than from the destination. Taking it from the destination is how a missing partition becomes invisible: it is not in the list, so it is not checked.

Compare sets, not sizes

The check that actually closes the question compares identifier sets. Both sides can be enumerated, so shipping a full identifier list is possible but usually unnecessary — hash each identifier to a fixed-width integer and combine the hashes with an order-independent operation. One number per side, and it differs if any identifier is missing, extra or changed.

import hashlib

def digest64(s: str) -> int:
    return int.from_bytes(hashlib.blake2b(s.encode(), digest_size=8).digest(), "big")

def set_signature(ids) -> tuple[int, int, int]:
    n = 0
    xor = 0
    total = 0
    for i in ids:
        h = digest64(i)
        xor ^= h
        total = (total + h) % (2 ** 61 - 1)
        n += 1
    return n, xor, total

Three numbers rather than one, deliberately. The XOR is order-independent and cheap but is blind to a duplicated identifier — any value XORed twice cancels. The modular sum is not: a duplicate changes it. And the count catches the case where both aggregates happen to agree. Together they fail on omission, addition, alteration and duplication, which is the full set of things a batch job does wrong.

Compute the signature per partition, not once globally, and store the results. A per-partition signature tells you where to look; a global one only tells you to look. When they disagree, then pull the identifier lists for that one partition and diff them properly.

The source side of the comparison must be the system of record — the database, the object store, the document table — and not the old vector index. Comparing the new index against the old index proves you copied the old index’s gaps faithfully.

Rows that finished and then went stale

A backfill over a live corpus has a moving target. A document edited after its chunk was re-embedded is now represented by a vector of its old text, and every count and set check above passes, because the identifier is present.

This is what source_version is for. The completeness predicate is not “a vector exists for every source row” but “a vector exists for every source row, produced by the current model, from the current version of that row”. In SQL against a pgvector table joined to the source table, that is one query:

SELECT c.partition,
       count(*) FILTER (WHERE v.id IS NULL)                        AS missing,
       count(*) FILTER (WHERE v.embed_model <> :target_model)       AS wrong_model,
       count(*) FILTER (WHERE v.source_version <> c.version)        AS stale,
       count(*) FILTER (WHERE v.embed_dim <> :target_dim)           AS wrong_dim
FROM chunks c
LEFT JOIN chunk_vectors v ON v.chunk_id = c.id
GROUP BY c.partition
ORDER BY c.partition;

Four numbers per partition, all of which must be zero. Against a vector database rather than a SQL store the same four are counts under metadata filters — a count with a filter on the model tag, a count of identifiers present in the source and absent from the index — but the predicate is identical and so is the standard: zero, per partition, not “small”.

If stale is non-zero and stays non-zero because the corpus is continuously edited, that is not an audit failure — it is the steady state, and the thing to assert is a bound on the age of the oldest stale row rather than a zero. RAG index freshness covers running that as an ongoing property instead of a one-off gate.

Vectors that exist but are wrong

A record can be present, correctly tagged, current, and still not be an embedding of anything. Error paths in batch jobs produce this: a caught exception that writes a zero vector to keep the pipeline moving, a default returned from a wrapper, a truncated response deserialised into a short array and padded.

Zero and near-zero vectors. Count records whose L2 norm is below a small epsilon. A legitimate embedding is essentially never the zero vector, and on a cosine index a zero vector produces an undefined or degenerate similarity.
Non-finite values. Count records containing NaN or infinity. These sometimes survive a write and then quietly corrupt ranking.
Duplicated vectors. Hash each vector’s bytes and count the most common hashes. A handful of duplicates is normal — identical chunks exist. Ten thousand records sharing one vector is a fallback value written by an exception handler, and it is the single most informative check on this list because it finds a class of bug nothing else does.
Dimension. Assert every vector has the expected width. Most stores enforce this, but a store with a lenient schema, or a collection created with the wrong dimension, will not.

Run these on a sample if the corpus is large enough that a full scan is expensive — a uniform sample of a hundred thousand records will find any systematic problem, and systematic is what these are.

Wiring the check into the cutover

Make the audit a job with an exit code, not a notebook. It runs from CI, on a schedule during the backfill, and on demand.
Have it emit one row per partition with all the counts above, and fail on any non-zero in the must-be-zero columns. Print the first twenty offending identifiers so the failure is actionable rather than just red.
Run it during the backfill, not only at the end. A partition that has been stuck at the same count for two hours is a worker that died, and you want to know at hour two rather than at the end.
Require it to pass twice, separated by a quiet period with the backfill stopped. One pass can coincide with in-flight writes; two passes with the same signature means the state has settled.
Only then flip reads, and keep the old index until the quality comparison in what changes in search quality after a migration has run against the frozen query set.
Keep the job afterwards, on a daily schedule. The same predicate that proves a migration finished proves that ingestion has not silently stopped, which is the same failure arriving six months later.