Describe the bug
We've been following this project for a while. We hit the problem below while testing replay_increment with a caching target.
replay_increment plans against the target's layout, flushes the cache, then writes each key at its source slot (batched_dynamicemb_tables.py:1934-1943). The problem: flush_cache can grow and rehash the backing store (_expand_storage_for_cache_flush_if_needed, key_value_table.py:3949), so the slots the plan just validated go stale and the write fails midway:
RuntimeError: replay_increment: 82 of 100 keys could not be written at their source slot on table 0 ... or the slot was held by another writer, meaning something ran against this table concurrently with the replay.
Nothing ran concurrently -- it was the flush itself.
Steps/Code to reproduce bug
L40. Both tables start at capacity 512; the flush inside replay_increment grows the target to 1152.
import torch
from dynamicemb import (DynamicEmbCheckMode, DynamicEmbPoolingMode,
DynamicEmbScoreStrategy, DynamicEmbTableOptions, EmbOptimType, EvictedItemMode)
from dynamicemb.batched_dynamicemb_tables import BatchedDynamicEmbeddingTablesV2
def make_model(dev):
opt = DynamicEmbTableOptions(
index_type=torch.int64, embedding_dtype=torch.float32, device_id=dev,
dim=8, max_capacity=4096, init_capacity=512, bucket_capacity=128,
safe_check_mode=DynamicEmbCheckMode.IGNORE, local_hbm_for_values=65536,
score_strategy=DynamicEmbScoreStrategy.TIMESTAMP, caching=True,
evicted_item_mode=EvictedItemMode.DISCARD)
return BatchedDynamicEmbeddingTablesV2(
table_options=[opt], output_dtype=torch.float32, table_names=["t_0"],
feature_table_map=[0], pooling_mode=DynamicEmbPoolingMode.SUM,
use_index_dedup=False, optimizer=EmbOptimType.SGD)
def touch(m, keys, device):
idx = torch.tensor(keys, dtype=torch.int64, device=device)
off = torch.arange(0, len(keys) + 1, dtype=torch.int64, device=device)
m(idx, off)
torch.cuda.synchronize()
device = torch.device("cuda:0")
src, dst = make_model(0), make_model(0)
delta_keys, extra_keys = list(range(1001, 1101)), list(range(5001, 5451))
touch(src, delta_keys, device)
delta = src.incremental_dump({"t_0": 0})
all_keys = delta_keys + extra_keys
for i in range(0, len(all_keys), 50): # small batches, or prefetch expands the table first
touch(dst, all_keys[i : i + 50], device)
dst.replay_increment(delta)Expected behavior
replay_increment either replays the delta or rejects it before writing anything -- the docstring promises ValueError "before this module writes anything" for layout mismatches. Here it fails midway and leaves a partial replay.
Environment details (please complete the following information):
Ubuntu 22.04, NVIDIA L40
Python 3.10, torch 2.10.0+cuda 12.8
Additional context
The fix is one line: re-run _plan_replay after the flush, so a changed layout rejects up front with the normal ValueError before anything is written. Verified: the repro rejects cleanly, the 26 existing replay tests pass, and a new regression test fails on the old code.
Easy to miss: with the same cache pressure on both sides, the source expands at its own dump first, so the target gets a clean rejection instead of a crash.
Can send a PR.
Thanks for the great work on this project!
By submitting this issue, you agree to follow our code of conduct and our contributing guidelines.
Describe the bug
We've been following this project for a while. We hit the problem below while testing replay_increment with a caching target.
replay_increment plans against the target's layout, flushes the cache, then writes each key at its source slot (batched_dynamicemb_tables.py:1934-1943). The problem: flush_cache can grow and rehash the backing store (
_expand_storage_for_cache_flush_if_needed, key_value_table.py:3949), so the slots the plan just validated go stale and the write fails midway:Nothing ran concurrently -- it was the flush itself.
Steps/Code to reproduce bug
L40. Both tables start at capacity 512; the flush inside replay_increment grows the target to 1152.
Expected behavior
replay_increment either replays the delta or rejects it before writing anything -- the docstring promises ValueError "before this module writes anything" for layout mismatches. Here it fails midway and leaves a partial replay.
Environment details (please complete the following information):
Ubuntu 22.04, NVIDIA L40
Python 3.10, torch 2.10.0+cuda 12.8
Additional context
The fix is one line: re-run
_plan_replayafter the flush, so a changed layout rejects up front with the normal ValueError before anything is written. Verified: the repro rejects cleanly, the 26 existing replay tests pass, and a new regression test fails on the old code.Easy to miss: with the same cache pressure on both sides, the source expands at its own dump first, so the target gets a clean rejection instead of a crash.
Can send a PR.
Thanks for the great work on this project!
By submitting this issue, you agree to follow our code of conduct and our contributing guidelines.