diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 86e11790..1ac96b07 100644 --- a/docs/docs/index/ivf.md +++ b/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. diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index f96c5d2f..113c443f 100644 --- a/docs/docs/index/usage.md +++ b/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. diff --git a/docs/docs/python_migration.md b/docs/docs/python_migration.md index 52fad117..c14a1bc6 100644 --- a/docs/docs/python_migration.md +++ b/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 diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 62d4f0b5..fe59efbd 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -252,7 +252,9 @@ 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 @@ -260,6 +262,11 @@ class IvfIndex: 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 ----- diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 779eae5d..ad2b76bb 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/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 num_threads) { const auto operation = access_.write(); auto centroids_array = ensure_2d_array(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,6 +617,8 @@ 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 @@ -624,6 +626,11 @@ 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 ----- diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 777b9823..32cbdbf9 100644 --- a/tests/python/test_ivf.py +++ b/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) diff --git a/tests/python/test_ivf_initializers.py b/tests/python/test_ivf_initializers.py index 96d8e1a9..1afdf8b0 100644 --- a/tests/python/test_ivf_initializers.py +++ b/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 diff --git a/tests/python/test_thread_counts.py b/tests/python/test_thread_counts.py index d3f5267a..a530523c 100644 --- a/tests/python/test_thread_counts.py +++ b/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 diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py index a0b27263..4d2a05cd 100644 --- a/tests/python/typing/public_api.py +++ b/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),))