feat(python): let IvfIndex.set_centroids take num_threads

docs/docs/index/ivf.md

@@ -223,7 +223,9 @@ index.add(vectors)
   keeping its default graph settings. The threshold is provisional and has not
   been measured.
 - `set_centroids(centroids)` uses centroids of shape `(num_clusters, dim)`
-  computed elsewhere instead of training.
+  computed elsewhere instead of training. Its keyword `num_threads` sets the
+  workers that rotate the centroids and build the routing; only the HNSW routing
+  graph varies with it.
 - Training or setting centroids again replaces them while the index stores no
   vectors; once it stores any (including removed ones), both raise
   `RuntimeError`. A failed call leaves the index unchanged.

docs/docs/index/usage.md

@@ -144,11 +144,13 @@ fewer workers. On macOS, OpenMP runtimes bundled by different packages (for
 example Faiss or PyTorch wheels) can crash when more than one of them starts
 worker threads in the same process; set `OMP_NUM_THREADS=1` before the first
 native call, or pass `num_threads=1`, in such processes.
-`train` and `add` on every index default to all CPUs (`None` or `0`); the
-clustering classes default to `0`. `search` defaults to one thread (`None` or `1`),
+`train` and `add` on every index, and IVF-RaBitQ `set_centroids`, default to all
+CPUs (`None` or `0`); the clustering classes default to `0`. `search` defaults to one thread (`None` or `1`),
 as do the deprecated `build`, which requires an integer, and `search_batch`.
-Multithreaded HNSW-RaBitQ `add` links points in a nondeterministic order, so pass
-`num_threads=1` when a build must be reproducible.
+Multithreaded HNSW-RaBitQ `add` links points in a nondeterministic order, and so
+does a multithreaded build of IVF-RaBitQ's HNSW centroid routing by `train`,
+`set_centroids` or the first `add`; pass `num_threads=1` when a build must be
+reproducible.
 On Linux, CPU detection accounts for affinity and OpenMP binding. Apply binding
 before starting the program.
 

docs/docs/python_migration.md

@@ -42,7 +42,9 @@ work if you want to run k-means yourself.
 
 - **Thread defaults.** `train` and `add` now default to every available CPU
   (`num_threads=None`, the same as `0`), capped by `OMP_NUM_THREADS` since 0.8.1; 0.7 `build` and
-  `add` defaulted to one thread. This applies to every `add` call that does not
+  `add` defaulted to one thread. Since 0.8.5, IVF-RaBitQ `set_centroids` follows
+  the same default and takes `num_threads`; before, it always used every CPU and
+  ignored `OMP_NUM_THREADS`. This applies to every `add` call that does not
   pass `num_threads`, including 0.7 code such as `index.add(x)`, without a
   warning. `search` still defaults to one thread, and the deprecated `build`
   keeps its default of 1.
@@ -51,6 +53,11 @@ work if you want to run k-means yourself.
   to `add` for the same graph on every run. With one thread and the same seed,
   `set_centroids` followed by `add` produces exactly the index that 0.7 `build`
   produced from the same centroids and cluster IDs.
+- **IVF-RaBitQ HNSW routing builds are not reproducible by default.** 0.7 `build`
+  built the routing graph with its own thread count, one by default. With more
+  than one thread, `set_centroids`, `train` and a training `add` build it in a
+  nondeterministic order; pass `num_threads=1` (to `set_centroids` since 0.8.5)
+  and a fixed `seed` for the same graph on every run.
 - **`k` above the point count.** 0.8.0 and 0.8.1 raised `ValueError` when `k`
   exceeded `num_stored_points` (HNSW-RaBitQ) or the point count (IVF-RaBitQ,
   SymphonyQG). Since 0.8.2, every index accepts any `k >= 1` and fills missing

python_bindings/__init__.pyi

@@ -252,14 +252,21 @@ class IvfIndex:
         raises RuntimeError afterwards. A failed call leaves the index unchanged.
         """
 
-    def set_centroids(self, centroids: npt.ArrayLike) -> None:
+    def set_centroids(
+        self, centroids: npt.ArrayLike, *, num_threads: SupportsIndex | None = None
+    ) -> None:
         """Use the given centroids instead of training.
 
         Parameters
         ----------
         centroids : array_like of float, shape (num_clusters, dim)
             Centroids in the original, unpadded coordinates. Converted to
             contiguous float32; a copy may be made.
+        num_threads : int, optional
+            Keyword-only. Worker count for rotating the centroids and building the
+            centroid routing. None (default) and 0 select the available CPU count,
+            capped by OMP_NUM_THREADS. Only the HNSW routing graph varies with it:
+            build it with 1 thread and a fixed seed for a reproducible graph.
 
         Notes
         -----

python_bindings/ivf_bindings.cpp

@@ -132,14 +132,14 @@ class IvfIndex {
         index_->train(rows, x_ptr, cp);
     }
 
-    void set_centroids(py::handle centroids) {
+    void set_centroids(py::handle centroids, std::optional<size_t> num_threads) {
         const auto operation = access_.write();
         auto centroids_array = ensure_2d_array<float>(centroids, "centroids");
         check_trainable();
         check_centroid_shape(centroids_array, index_->num_clusters(), index_->dimension());
         const auto* centroids_ptr = centroids_array.data();
         py::gil_scoped_release release;
-        index_->set_centroids(centroids_ptr);
+        index_->set_centroids(centroids_ptr, num_threads.value_or(0));
     }
 
     // Deprecated: the 0.7 build. A point-count index is constructed; any other
@@ -617,13 +617,20 @@ raises RuntimeError afterwards. A failed call leaves the index unchanged.)doc"
             "set_centroids",
             &IvfIndex::set_centroids,
             py::arg("centroids"),
+            py::kw_only(),
+            py::arg("num_threads") = py::none(),
             R"doc(Use the given centroids instead of training.
 
 Parameters
 ----------
 centroids : array_like of float, shape (num_clusters, dim)
     Centroids in the original, unpadded coordinates. Converted to
     contiguous float32; a copy may be made.
+num_threads : int, optional
+    Keyword-only. Worker count for rotating the centroids and building the
+    centroid routing. None (default) and 0 select the available CPU count,
+    capped by OMP_NUM_THREADS. Only the HNSW routing graph varies with it:
+    build it with 1 thread and a fixed seed for a reproducible graph.
 
 Notes
 -----

tests/python/test_ivf.py

@@ -516,17 +516,16 @@ def test_hnsw_centroid_routing_roundtrip(tmp_path):
     count, dim = 128, 64
     centroids = rng.standard_normal((count, dim)).astype(np.float32)
     labels = np.arange(count, dtype=np.uint32)
-    idx = _built(
-        centroids, centroids, labels, nbits=1, num_threads=4, initializer="hnsw"
-    )
-    # set_centroids builds the routing graph with every CPU, and a parallel build may
-    # leave a centroid unreachable: check every centroid's self-recall in aggregate.
-    every, every_distances = idx.search(centroids, k=1, nprobe=4, num_threads=4)
-    found = every[:, 0] == labels
-    assert np.count_nonzero(found) >= 0.97 * count
-    np.testing.assert_allclose(every_distances[found], 0, rtol=0, atol=1e-5)
-    queries = centroids[[0, 17, 63, count - 1]]
+    idx = IvfIndex(dim, count, nbits=1, initializer="hnsw", seed=7)
+    # A seeded single-thread routing build gives the same graph on every run; a
+    # parallel build may leave a centroid unreachable (#216).
+    idx.set_centroids(centroids, num_threads=1)
+    idx.add(centroids, cluster_ids=labels, num_threads=4)
+    selected = np.array([0, 17, 63, count - 1])
+    queries = centroids[selected]
     ids, distances = idx.search(queries, k=1, nprobe=4, num_threads=4)
+    np.testing.assert_array_equal(ids[:, 0], selected)
+    np.testing.assert_allclose(distances, 0, rtol=0, atol=1e-5)
 
     path = str(tmp_path / "hnsw_centroid.index")
     idx.save(path)

tests/python/test_ivf_initializers.py

@@ -16,9 +16,7 @@ def test_hnsw_parallel_build_routes_and_reloads(tmp_path, metric, threads):
     data = np.random.default_rng(42).standard_normal((count, dim)).astype(np.float32)
     data /= np.linalg.norm(data, axis=1, keepdims=True)
     index = IvfIndex(dim, count, nbits=4, metric=metric, initializer="hnsw")
-    # set_centroids builds the routing graph with every CPU; `threads` applies to
-    # add and search.
-    index.set_centroids(data)
+    index.set_centroids(data, num_threads=threads)
     index.add(data, cluster_ids=np.arange(count, dtype=np.uint32), num_threads=threads)
     # A parallel build may leave a centroid unreachable in the routing graph, so
     # check the aggregate self-recall of every point (one point per cluster).
@@ -37,6 +35,24 @@ def test_hnsw_parallel_build_routes_and_reloads(tmp_path, metric, threads):
         np.testing.assert_array_equal(actual, reference)
 
 
+@pytest.mark.parametrize("metric", ["l2", "ip"])
+def test_single_thread_set_centroids_builds_reproducible_hnsw_routing(tmp_path, metric):
+    count, dim = 1024, 64
+    data = np.random.default_rng(42).standard_normal((count, dim)).astype(np.float32)
+    saved = []
+    for run in range(2):
+        index = IvfIndex(dim, count, nbits=1, metric=metric, initializer="hnsw", seed=7)
+        index.set_centroids(data, num_threads=1)
+        index.add(data, cluster_ids=np.arange(count, dtype=np.uint32), num_threads=4)
+        path = tmp_path / f"run-{run}.index"
+        index.save(str(path))
+        sidecar = Path(str(path) + ".hnsw")  # absent once the file embeds routing
+        saved.append(
+            path.read_bytes() + (sidecar.read_bytes() if sidecar.exists() else b"")
+        )
+    assert saved[0] == saved[1]
+
+
 @pytest.mark.parametrize("metric", ["l2", "ip"])
 def test_hnsw_concurrent_routing_windows_match_serial(tmp_path, metric):
     count, dim = 1024, 64

tests/python/test_thread_counts.py

@@ -83,3 +83,35 @@ def test_automatic_thread_count_respects_omp_num_threads(limit):
         timeout=30,
     )
     assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_default_set_centroids_respects_omp_num_threads(tmp_path):
+    # IVF set_centroids used to build HNSW routing with every CPU, whatever the
+    # OpenMP limit. With one OpenMP thread, two default builds must agree.
+    code = """
+import sys
+from pathlib import Path
+import numpy as np
+from rabitqlib import IvfIndex
+data = np.random.default_rng(42).standard_normal((1024, 64)).astype(np.float32)
+saved = []
+for run in range(2):
+    index = IvfIndex(64, 1024, nbits=1, initializer="hnsw", seed=7)
+    index.set_centroids(data)
+    index.add(data[:1], cluster_ids=[0], num_threads=1)
+    path = Path(sys.argv[1]) / f"run-{run}.index"
+    index.save(str(path))
+    sidecar = Path(str(path) + ".hnsw")
+    saved.append(path.read_bytes() + (sidecar.read_bytes() if sidecar.exists() else b""))
+assert saved[0] == saved[1]
+"""
+    env = os.environ.copy()
+    env["OMP_NUM_THREADS"] = "1"
+    result = subprocess.run(
+        [sys.executable, "-c", code, str(tmp_path)],
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr

tests/python/typing/public_api.py

@@ -205,7 +205,7 @@ def check_input_contracts() -> None:
         IvfIndex(64, 1, nbits=4, initializer=initializer)
     for ids in (signed, unsigned, [0] * 80, (0,) * 80):
         ivf = IvfIndex(64, 1, nbits=4)
-        ivf.set_centroids(centroids)
+        ivf.set_centroids(centroids, num_threads=1)
         ivf.add(x, cluster_ids=ids)
         ivf.add(x[:1], cluster_ids=[np.int64(0)])
         ivf.add(x[:1], cluster_ids=(np.uint32(0),))