[feat] tombstone evicted dynamic embedding keys in delta dump - #638
[feat] tombstone evicted dynamic embedding keys in delta dump#638eric-gecheng wants to merge 1 commit into
Conversation
Evicted dynamicemb keys were silently skipped at dump time, so their stale FeatureStore rows kept serving forever and the processor could never reclaim the memory. The dump now drains pop_evicted_keys after the tracker pass (capturing flush-induced evictions in the same window), subtracts keys already republished as real rows, and uploads the rest as all-zero tombstones through the existing chunk/quant path; the polling processor detects the zeros and deletes the key from NvEmbeddings. The new dump_evicted_tombstones switch (default true) gates the pass and auto-arms dynamicemb RETAIN_KEY recording for exactly those runs, degrading to a one-time warning on old dynamicemb builds. Also fix the trailing-comma JSON error in .pyre_configuration that broke pyre startup.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| ), | ||
| table_fqn=fqn, | ||
| key_ids=evicted, | ||
| embeddings=torch.zeros((evicted.numel(), emb_dim), dtype=torch.float32), |
There was a problem hiding this comment.
This tombstone is indistinguishable from a valid embedding after upload drops the source column. Dynamic embeddings support CONSTANT initialization with the default value 0.0 (and frozen rows can remain zero), so the Processor can delete a live key. Please use an explicit delete marker in the persisted wire contract, or otherwise enforce an end-to-end representation that real rows cannot produce.
| ), | ||
| table_fqn=fqn, | ||
| key_ids=evicted, | ||
| embeddings=torch.zeros((evicted.numel(), emb_dim), dtype=torch.float32), |
There was a problem hiding this comment.
This expands every evicted key into a dense E × embedding_dim FP32 tensor and retains it in the Arrow chunks/upload queue. For example, 1M evictions at dimension 128 require roughly 512 MB before metadata and other copies, synchronously on the training worker. Please emit tombstones in bounded chunks or use a compact delete representation.
| enable_delta_embedding_dump = train_config.HasField("delta_embedding_dump_config") | ||
| # Arm evicted-key retention before feature/model building so the dump's | ||
| # pop_evicted_keys drain has a buffer to consume on every dynamicemb table. | ||
| dynamicemb_util.set_auto_retain_evicted_keys( |
There was a problem hiding this comment.
set_auto_retain_evicted_keys() mutates process-global state, but this path never restores it. A later export/predict/model build in the same process can inherit RETAIN_KEY without constructing a dumper to drain it, so the eviction buffer can grow indefinitely. Scope/reset the flag after planning/sharding (including exception paths), or pass the option explicitly.
| // enabled, delta dump mode also arms dynamicemb's evicted-key recording | ||
| // (evicted_item_mode=RETAIN_KEY) automatically. Requires a dynamicemb | ||
| // build with EvictedItemMode; older builds degrade to a warning. | ||
| optional bool dump_evicted_tombstones = 7 [default = true]; |
There was a problem hiding this comment.
Defaulting this to true silently enables RETAIN_KEY for every existing delta-dump configuration. The retained-key buffer grows with eviction volume until the next dump, so high-churn tables or long intervals can OOM before the first drain. Please make this opt-in, or bound/account for the buffer independently of the dump interval.
| # The tiny initial table capacity evicts keys during training; they | ||
| # must reach the shards as tombstones or the pop_evicted_keys drain | ||
| # was never exercised. | ||
| self.assertTrue( |
There was a problem hiding this comment.
This test does not deterministically force an eviction: init_capacity_per_rank = 128 is an initial growth allocation, while the fixture's max_capacity values are 1,000,000 and 100,000, above the generated unique-ID counts. Configure a deliberately small max_capacity with known IDs before requiring a tombstone. Also gate this assertion with the new capability check, since production explicitly supports older builds without pop_evicted_keys.
|
Static review complete; no tests or builds were run. I left five inline comments covering zero-vector tombstone ambiguity, process-global retention state, retained-key memory growth, dense tombstone allocation, and deterministic compatibility-aware eviction coverage. One general documentation gap: please add Chinese user-facing docs for the default-on switch and opt-out, required dynamicemb and Processor versions, retained-key memory behavior, and INT8 tombstone encoding. No additional security-specific issues stood out. |
1 participant
Evicted dynamicemb keys were silently skipped at dump time, so their stale FeatureStore rows kept serving forever and the processor could never reclaim the memory. The dump now drains pop_evicted_keys after the tracker pass (capturing flush-induced evictions in the same window), subtracts keys already republished as real rows, and uploads the rest as all-zero tombstones through the existing chunk/quant path; the polling processor detects the zeros and deletes the key from NvEmbeddings. The new dump_evicted_tombstones switch (default true) gates the pass and auto-arms dynamicemb RETAIN_KEY recording for exactly those runs, degrading to a one-time warning on old dynamicemb builds. Also fix the trailing-comma JSON error in .pyre_configuration that broke pyre startup.