fix: harden deterministic IVF-PQ training

This commit is contained in:
Gatefixer
2026-08-06 06:32:58 +00:00
parent cc321e9801
commit 845868a343
8 changed files with 585 additions and 145 deletions
+9 -1
View File
@@ -772,7 +772,9 @@ class IvfPq:
seed: int, optional
Seed used for deterministic sampling and training. Given identical data in
the same row order and identical index parameters, the same seed produces
the same IVF centroids and PQ codebook. If omitted, training remains random.
the same IVF centroids and PQ codebook. This option is supported for local
CPU index builds; remote and accelerator builds reject it explicitly. If
omitted, training remains random.
target_partition_size: int, default is 8192
@@ -794,6 +796,12 @@ class IvfPq:
# create_index() dispatches to pylance to build the index on the accelerator.
accelerator: Optional[str] = None
def __post_init__(self):
if self.seed is not None and self.accelerator is not None:
raise ValueError(
"IvfPq seed is not supported with accelerator-based index training"
)
@dataclass
class IvfRq:
+5
View File
@@ -395,6 +395,11 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
assert stats.num_indices == 1
def test_seeded_ivfpq_rejects_accelerator():
with pytest.raises(ValueError, match="seed is not supported with accelerator"):
IvfPq(seed=42, accelerator="cuda")
@pytest.mark.asyncio
async def test_create_ivfrq_index(some_table: AsyncTable):
await some_table.create_index("vector", config=IvfRq(num_bits=1))