diff --git a/CMakeLists.txt b/CMakeLists.txt index ac9eac53..45ac186a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,6 @@ cmake_minimum_required(VERSION 3.20) -project(RaBitQLib VERSION 0.8.0 LANGUAGES CXX) +project(RaBitQLib VERSION 0.8.1 LANGUAGES CXX) include(CMakePackageConfigHelpers) include(GNUInstallDirs) diff --git a/README.md b/README.md index d247f358..a730983a 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,8 @@ approximate nearest-neighbor search. See the [research papers](https://vectordb- ## News +- **v0.8.1:** The default `num_threads=0` now respects `OMP_NUM_THREADS`, fixing a + macOS crash when Faiss and PyTorch share the process. See the [release notes](https://vectordb-ntu.github.io/RaBitQ-Library/releases/0.8.1/). - **v0.8.0:** Faiss-style `train`, `add` and `search` for all Python indexes, with 0.7 calls kept as deprecated forms until 0.9; optional k-means sampling and assignment; stricter input validation and atomic index saves. C++ query-factor @@ -52,7 +54,7 @@ settings, and limitations. ### Python ```bash -python -m pip install --upgrade "rabitqlib>=0.8.0" +python -m pip install --upgrade "rabitqlib>=0.8.1" ``` Wheels support CPython 3.11–3.14 on Linux x86-64/ARM64, Windows x86-64, diff --git a/docs/docs/clustering.md b/docs/docs/clustering.md index 9aabbc5c..40e0bc8e 100644 --- a/docs/docs/clustering.md +++ b/docs/docs/clustering.md @@ -95,7 +95,7 @@ Set Python keyword arguments or fields of the corresponding C++ parameter type: | --- | --- | --- | | `niter` | `25` | Maximum clustering iterations. | | `early_stop_threshold` | `0.0` | Stop from iteration 2 when the absolute relative objective change is at most this value; range `[0, 1]`. Zero stops only on an unchanged objective. | -| `num_threads` | `0` | Use detected available logical CPUs; positive requests are capped at that count. See [threading](index/usage.md#threading-and-file-paths). | +| `num_threads` | `0` | Use detected available logical CPUs, capped by `OMP_NUM_THREADS`; positive requests are capped at the CPU count. See [threading](index/usage.md#threading-and-file-paths). | | `spherical` | `false` | Normalize centroids and assign by inner product; input must already be normalized. | | `seed` | `42` | Seed for centroid selection and rotation. | | `min_points_per_centroid` | `39` | Warn below this training-points-per-centroid ratio when verbose. | diff --git a/docs/docs/cpp_integration.md b/docs/docs/cpp_integration.md index b99764e9..f68be782 100644 --- a/docs/docs/cpp_integration.md +++ b/docs/docs/cpp_integration.md @@ -58,12 +58,12 @@ The core example is built as the `core_example` CMake target. For the GIST bench The [core quantization v1 API](rabitq/core_api.md) has a C++ source-compatibility contract. Other C++ APIs may evolve, and no cross-version C++ binary ABI is promised. -Pin the `v0.8.0` release and include RaBitQ-Library as a Git submodule: +Pin the `v0.8.1` release and include RaBitQ-Library as a Git submodule: ```bash git submodule add https://github.com/VectorDB-NTU/RaBitQ-Library.git third_party/rabitqlib git submodule update --init --recursive -git -C third_party/rabitqlib checkout v0.8.0 +git -C third_party/rabitqlib checkout v0.8.1 git add third_party/rabitqlib ``` diff --git a/docs/docs/index.md b/docs/docs/index.md index eae6276e..9a2cdae4 100644 --- a/docs/docs/index.md +++ b/docs/docs/index.md @@ -26,7 +26,7 @@ estimation into your own C++ search engine.

```bash -python -m pip install "rabitqlib>=0.8.0" +python -m pip install "rabitqlib>=0.8.1" ``` [Installation and supported platforms](installation.md) · @@ -112,5 +112,5 @@ on your own dataset when choosing an index. ## Explore further - [Algorithm overview](rabitq/rabitq.md) for the rotation, encoding, and scoring workflow. -- [RaBitQ Library 0.8.0 release notes](releases/0.8.0.md) for the latest changes and upgrade guidance. +- [RaBitQ Library 0.8.1 release notes](releases/0.8.1.md) for the latest changes and upgrade guidance. - [Papers and citations](papers.md) for the original RaBitQ method, its multi-bit extension, and SymphonyQG. diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index 1bd2b06a..d7cb4e69 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -137,8 +137,13 @@ thread. Returned result arrays own their storage. These access checks belong to the Python wrappers. C++ callers must synchronize index updates themselves and use separate search output buffers. -`num_threads=0` selects the detected available logical CPU count. Positive values -set an upper limit, capped at that count; small workloads may use fewer workers. +`num_threads=0` selects the detected available logical CPU count, capped by the +OpenMP thread limit (`OMP_NUM_THREADS`, or `omp_set_num_threads` in C++). Positive +values set an upper limit, capped at the CPU count only; small workloads may use +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 `0`, as does clustering; `search` and the deprecated `build` default to one thread. Multithreaded HNSW-RaBitQ `add` links points in a nondeterministic order, so pass diff --git a/docs/docs/installation.md b/docs/docs/installation.md index 16d41570..2b8e879a 100644 --- a/docs/docs/installation.md +++ b/docs/docs/installation.md @@ -1,6 +1,6 @@ # Installation -Platform requirements and source installation for RaBitQ Library **0.8.0**. +Platform requirements and source installation for RaBitQ Library **0.8.1**. For the first Python example, see [Quick Start](quick_start.md). ## Requirements @@ -40,7 +40,7 @@ build can use scalar kernels without the FastScan CPU baseline. ## Python wheels ```bash -python -m pip install "rabitqlib>=0.8.0" +python -m pip install "rabitqlib>=0.8.1" ``` For changes from 0.6.0, see the [0.7.0 upgrade notes](releases/0.7.0.md). @@ -61,11 +61,11 @@ importing those classes. Array shapes and parameter ranges are still checked at ## Get the source -For Python source builds or C++ use, obtain the 0.8.0 checkout once. With Git +For Python source builds or C++ use, obtain the 0.8.1 checkout once. With Git installed, run: ```bash -git clone --branch v0.8.0 --depth 1 https://github.com/VectorDB-NTU/RaBitQ-Library.git +git clone --branch v0.8.1 --depth 1 https://github.com/VectorDB-NTU/RaBitQ-Library.git cd RaBitQ-Library ``` diff --git a/docs/docs/python_migration.md b/docs/docs/python_migration.md index da1bc83d..baf942ce 100644 --- a/docs/docs/python_migration.md +++ b/docs/docs/python_migration.md @@ -40,10 +40,11 @@ work if you want to run k-means yourself. ### Behavior changes - **Thread defaults.** `train` and `add` now default to every available CPU - (`num_threads=0`); 0.7 `build` and `add` defaulted to one thread. 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. + (`num_threads=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 + 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. - **HNSW-RaBitQ builds are not reproducible by default.** With more than one thread, `add` links points in a nondeterministic order. Pass `num_threads=1` to `add` for the same graph on every run. With one thread and the same seed, diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index bbaad793..817a3dbc 100644 --- a/docs/docs/quick_start.md +++ b/docs/docs/quick_start.md @@ -1,14 +1,14 @@ # Quick start Build, search, save and reload a small IVF-RaBitQ index with Python. This guide describes -**0.8.0** and uses synthetic data, so no dataset download is required. Coming from 0.7? +**0.8.1** and uses synthetic data, so no dataset download is required. Coming from 0.7? See the [migration guide](python_migration.md). For C++, go directly to [C++ integration](cpp_integration.md). ## Install ```bash -python -m pip install "rabitqlib>=0.8.0" +python -m pip install "rabitqlib>=0.8.1" ``` Wheels support CPython 3.11–3.14 on Linux x86-64/ARM64, Windows x86-64 and @@ -47,9 +47,9 @@ print(ids[0]) iterations on at most 256 sampled vectors per centroid, spherical k-means for `metric="ip"`, and the index `seed`. Change `index.cp` before training to tune them, or call `index.set_centroids(centroids)` to use centroids you computed -elsewhere. `train` and `add` use every available CPU by default; pass -`num_threads=` to limit them. Without `seed`, each index draws a random rotation -seed and reports it as `index.seed`. +elsewhere. `train` and `add` use every available CPU by default, capped by +`OMP_NUM_THREADS`; pass `num_threads=` to limit them. Without `seed`, each index +draws a random rotation seed and reports it as `index.seed`. ## Read the results diff --git a/docs/docs/releases/0.8.1.md b/docs/docs/releases/0.8.1.md new file mode 100644 index 00000000..8077480d --- /dev/null +++ b/docs/docs/releases/0.8.1.md @@ -0,0 +1,32 @@ +# RaBitQ Library 0.8.1 + +This patch release fixes a crash on macOS when RaBitQ Library runs in the same +process as Faiss and PyTorch, reported through txtai +([#207](https://github.com/VectorDB-NTU/RaBitQ-Library/issues/207)). + +## Fixes + +- `num_threads=0`, the default for `train`, `add` and clustering, now selects + the available logical CPUs capped by the OpenMP thread limit + (`OMP_NUM_THREADS`, or `omp_set_num_threads` in C++), as Faiss does. Positive + requests are unchanged and remain capped only by the CPU count. Without + `OMP_NUM_THREADS`, the default still uses every available CPU. +- In 0.8.0, 0.7-style `add(x)` calls switched to every CPU and ignored + `OMP_NUM_THREADS`. On macOS, the Faiss, PyTorch and RaBitQ Library wheels each + bundle their own OpenMP runtime; with Faiss and PyTorch loaded, starting RaBitQ + Library's worker threads crashed with `OMP: Error #179` followed by a + segmentation fault. Applications that set `OMP_NUM_THREADS=1` for this reason, + such as txtai, run `add` single-threaded again. Processes that mix these + runtimes without that setting should set it before the first native call or + pass `num_threads=1`; see [threading](../index/usage.md#threading-and-file-paths). + +## Upgrading + +From 0.8.0, public APIs, index file formats and the stable core API v1 are +unchanged; no index conversion is required. If you set `OMP_NUM_THREADS` and want +`train` or `add` to use more threads, pass `num_threads` explicitly. Rebuild C++ +consumers with matching headers and libraries; no cross-version C++ binary ABI is +promised. + +When upgrading from 0.7, follow the [0.8.0 release notes](0.8.0.md) and the +[migration guide](../python_migration.md). diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 753e30b2..f8e936ab 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -13,6 +13,7 @@ nav: - Installation: installation.md - C++ Integration: cpp_integration.md - Release notes: + - "0.8.1": releases/0.8.1.md - "0.8.0": releases/0.8.0.md - "0.7.2": releases/0.7.2.md - "0.7.1": releases/0.7.1.md diff --git a/include/rabitqlib/utils/tools.hpp b/include/rabitqlib/utils/tools.hpp index af34185b..11ea3b29 100644 --- a/include/rabitqlib/utils/tools.hpp +++ b/include/rabitqlib/utils/tools.hpp @@ -85,13 +85,19 @@ inline bool is_powerof2(size_t n) { return n > 0 && (n & (n - 1)) == 0; } // Available logical CPUs in the calling context, with a hardware-count fallback. size_t total_threads(); -// Zero selects all available CPUs; explicit requests cannot exceed that count. +// total_threads() capped by the OpenMP thread limit (OMP_NUM_THREADS or +// omp_set_num_threads), so applications can limit every automatic team. +size_t default_threads(); + +// Zero selects default_threads(); explicit requests cannot exceed total_threads(). inline size_t resolve_num_threads(size_t requested) { if (requested == 1) { return 1; } - const size_t available = total_threads(); - return requested == 0 ? available : std::min(requested, available); + if (requested == 0) { + return default_threads(); + } + return std::min(requested, total_threads()); } namespace detail { diff --git a/pyproject.toml b/pyproject.toml index 80094eff..3ab63ef1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "rabitqlib" -version = "0.8.0" +version = "0.8.1" description = "Official RaBitQ library for vector quantization and approximate nearest-neighbor search" readme = "README.md" requires-python = ">=3.11" diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index 017c7252..9ce77a7f 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -229,7 +229,8 @@ def build( Centroid ID per vector, in [0, num_clusters). Floating-point and boolean IDs are rejected. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. fast_quantization : bool, optional Use faster quantizer configuration. Default is False. """ @@ -328,7 +329,8 @@ def build(self, data, ef_construction, num_threads=1, init="pipnn"): ef_construction : int Positive construction search breadth fitting uint32. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. init : {'pipnn', 'random'}, optional Graph initialization method. Default is 'pipnn'. """ @@ -522,7 +524,8 @@ class QGKMeans(_KMeans): Stop when the absolute relative objective change is at most this threshold. Default is 0.0. num_threads : int, optional - Worker count; 0 selects the available CPU count. Default is 0. + Worker count; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Default is 0. final_assignment : FinalAssignmentMode, optional Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is @@ -586,7 +589,8 @@ class RaBitQKMeans(_KMeans): Stop when the absolute relative objective change is at most this threshold. Default is 0.0. num_threads : int, optional - Worker count; 0 selects the available CPU count. Default is 0. + Worker count; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Default is 0. final_assignment : FinalAssignmentMode, optional Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index ca545545..629fdf81 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -237,7 +237,8 @@ class IvfIndex: n >= num_clusters. Converted to contiguous float32; a copy may be made. num_threads : int, optional Keyword-only. Worker count for this call. None (default) uses - cp.num_threads; 0 selects the available CPU count. + cp.num_threads; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Notes ----- @@ -294,7 +295,8 @@ class IvfIndex: Centroid ID per vector, in [0, num_clusters). Floating-point and boolean IDs are rejected. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. fast_quantization : bool, optional Use faster quantizer configuration. Default is False. """ @@ -390,7 +392,8 @@ class IvfIndex: Enable HACC. None (default) selects it for 4-9 bits and standard FastScan otherwise. Raw storage uses standard FastScan. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Returns ------- @@ -648,7 +651,8 @@ class HnswIndex: n >= num_clusters. Converted to contiguous float32; a copy may be made. num_threads : int, optional Keyword-only. Worker count for this call. None (default) uses - cp.num_threads; 0 selects the available CPU count. + cp.num_threads; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Notes ----- @@ -704,7 +708,8 @@ class HnswIndex: Centroid ID per vector, in [0, num_clusters). Floating-point and boolean IDs are rejected. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. fast_quantization : bool, optional Use faster quantizer configuration. Default is False. """ @@ -798,7 +803,8 @@ class HnswIndex: Search breadth. None or 0 (default) selects max(k, 10); values above the stored point count are clamped. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Returns ------- @@ -1111,7 +1117,8 @@ class SymqgIndex: ef_construction : int Positive construction search breadth fitting uint32. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. init : {'pipnn', 'random'}, optional Graph initialization method. Default is 'pipnn'. """ @@ -1136,7 +1143,8 @@ class SymqgIndex: Search breadth. None or 0 (default) selects max(k, 10); smaller values are raised to k. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Returns ------- @@ -1222,7 +1230,8 @@ class QGKMeans: Stop when the absolute relative objective change is at most this threshold. Default is 0.0. num_threads : int, optional - Worker count; 0 selects the available CPU count. Default is 0. + Worker count; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Default is 0. final_assignment : FinalAssignmentMode, optional Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is @@ -1306,7 +1315,8 @@ class QGKMeans: Stop when the absolute relative objective change is at most this threshold. Default is 0.0. num_threads : int, optional - Worker count; 0 selects the available CPU count. Default is 0. + Worker count; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Default is 0. final_assignment : FinalAssignmentMode, optional Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is @@ -1464,7 +1474,7 @@ class QGKMeansParameters: @property def num_threads(self) -> int: - """Worker count; 0 selects the available CPU count. Default is 0.""" + """Worker count; 0 selects the available CPU count, capped by OMP_NUM_THREADS. Default is 0.""" @num_threads.setter def num_threads(self, value: SupportsIndex) -> None: ... @@ -1528,7 +1538,8 @@ class RaBitQKMeans: Stop when the absolute relative objective change is at most this threshold. Default is 0.0. num_threads : int, optional - Worker count; 0 selects the available CPU count. Default is 0. + Worker count; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Default is 0. final_assignment : FinalAssignmentMode, optional Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is @@ -1604,7 +1615,8 @@ class RaBitQKMeans: Stop when the absolute relative objective change is at most this threshold. Default is 0.0. num_threads : int, optional - Worker count; 0 selects the available CPU count. Default is 0. + Worker count; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Default is 0. final_assignment : FinalAssignmentMode, optional Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is @@ -1760,7 +1772,7 @@ class RaBitQKMeansParameters: @property def num_threads(self) -> int: - """Worker count; 0 selects the available CPU count. Default is 0.""" + """Worker count; 0 selects the available CPU count, capped by OMP_NUM_THREADS. Default is 0.""" @num_threads.setter def num_threads(self, value: SupportsIndex) -> None: ... diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index ae8182bd..8327b4f1 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -537,7 +537,7 @@ x : array_like of float, shape (n, dim) n >= num_clusters. Converted to contiguous float32; a copy may be made. num_threads : int, optional Keyword-only. Worker count for this call. None (default) uses - cp.num_threads; 0 selects the available CPU count. + cp.num_threads; 0 selects the available CPU count, capped by OMP_NUM_THREADS. Notes ----- @@ -593,7 +593,8 @@ ef : int, optional Search breadth. None or 0 (default) selects max(k, 10); values above the stored point count are clamped. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Returns ------- diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index aa04a1d9..0a73b215 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -563,7 +563,7 @@ x : array_like of float, shape (n, dim) n >= num_clusters. Converted to contiguous float32; a copy may be made. num_threads : int, optional Keyword-only. Worker count for this call. None (default) uses - cp.num_threads; 0 selects the available CPU count. + cp.num_threads; 0 selects the available CPU count, capped by OMP_NUM_THREADS. Notes ----- @@ -672,7 +672,8 @@ high_accuracy : bool, optional Enable HACC. None (default) selects it for 4-9 bits and standard FastScan otherwise. Raw storage uses standard FastScan. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Returns ------- diff --git a/python_bindings/kmeans_bindings.cpp b/python_bindings/kmeans_bindings.cpp index 96077a8a..16574b7d 100644 --- a/python_bindings/kmeans_bindings.cpp +++ b/python_bindings/kmeans_bindings.cpp @@ -108,7 +108,7 @@ void bind_training_parameters(py::class_& binding) { .def_readwrite( "num_threads", &Parameters::num_threads, - R"doc(Worker count; 0 selects the available CPU count. Default is 0.)doc" + R"doc(Worker count; 0 selects the available CPU count, capped by OMP_NUM_THREADS. Default is 0.)doc" ) .def_readwrite( "final_assignment", diff --git a/python_bindings/rabitq_bindings.cpp b/python_bindings/rabitq_bindings.cpp index e4d0eb08..2b47b2ca 100644 --- a/python_bindings/rabitq_bindings.cpp +++ b/python_bindings/rabitq_bindings.cpp @@ -17,6 +17,7 @@ PYBIND11_MODULE(_rabitqlib, m) { m.doc() = "RabitQ Python bindings combined module"; m.def("_available_cpu_count", &rabitqlib::total_threads); + m.def("_default_thread_count", &rabitqlib::default_threads); // Register shared enums once py::enum_(m, "MetricType") diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index a91e7262..843131b3 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -381,7 +381,8 @@ ef : int, optional Search breadth. None or 0 (default) selects max(k, 10); smaller values are raised to k. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Worker count. Default is 1; 0 selects the available CPU count, + capped by OMP_NUM_THREADS. Returns ------- diff --git a/sample/python/README.md b/sample/python/README.md index 532caaf0..8806618b 100644 --- a/sample/python/README.md +++ b/sample/python/README.md @@ -1,7 +1,7 @@ # Python examples Run these commands from the repository root in your project Python environment -with `rabitqlib>=0.8.0` and NumPy installed. When working from source, +with `rabitqlib>=0.8.1` and NumPy installed. When working from source, [install the library from the same checkout](../../.github/CONTRIBUTING.md#python-changes) so the examples and native extension use matching APIs. Clustering and indexing both use RaBitQ Library; FAISS is needed only for `compare_with_faiss.py`. @@ -43,9 +43,10 @@ behavior. For inner product, pass `--metric ip` to every stage and provide normalized vectors. Use queries and ground truth prepared for the same metric. -`--num-threads 0` uses the library's available logical CPU count; each stage accepts -its own limit. See [threading](../../docs/docs/index/usage.md#threading-and-file-paths) -for affinity and OpenMP placement behavior. +`--num-threads 0` uses the library's available logical CPU count, capped by +`OMP_NUM_THREADS`; each stage accepts its own limit. See +[threading](../../docs/docs/index/usage.md#threading-and-file-paths) for affinity +and OpenMP placement behavior. ## HNSW-RaBitQ diff --git a/sample/python/kmeans_clustering.py b/sample/python/kmeans_clustering.py index 367ec8e7..79b99e6d 100644 --- a/sample/python/kmeans_clustering.py +++ b/sample/python/kmeans_clustering.py @@ -5,7 +5,7 @@ import numpy as np from rabitqlib import FinalAssignmentMode, QGKMeans, RaBitQKMeans -from rabitqlib._rabitqlib import _available_cpu_count +from rabitqlib._rabitqlib import _available_cpu_count, _default_thread_count from utils import read_fvecs @@ -26,9 +26,10 @@ def cluster_data( if method not in ("rabitq", "qg"): raise ValueError("method must be rabitq or qg") clustering_type = RaBitQKMeans if method == "rabitq" else QGKMeans - available_threads = _available_cpu_count() threads = ( - available_threads if num_threads == 0 else min(num_threads, available_threads) + _default_thread_count() + if num_threads == 0 + else min(num_threads, _available_cpu_count()) ) print(f"Clustering metric: {metric.upper()}, threads: {threads}") kmeans = clustering_type( diff --git a/src/utils/thread_count.cpp b/src/utils/thread_count.cpp index e2bc56e1..f4db721e 100644 --- a/src/utils/thread_count.cpp +++ b/src/utils/thread_count.cpp @@ -63,4 +63,15 @@ size_t total_threads() { return threads == 0 ? 1 : threads; } +size_t default_threads() { + const size_t available = total_threads(); +#if RABITQ_USE_OPENMP && defined(_OPENMP) + const int limit = omp_get_max_threads(); + if (limit > 0) { + return std::min(available, static_cast(limit)); + } +#endif + return available; +} + } // namespace rabitqlib diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index 27c19020..35c0c436 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -57,6 +57,7 @@ def test_clustering_example_uses_native_cpu_budget(monkeypatch, requested, expec module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) monkeypatch.setattr(module, "_available_cpu_count", lambda: 2) + monkeypatch.setattr(module, "_default_thread_count", lambda: 2) calls = [] def clustering(*args, **kwargs): diff --git a/tests/python/test_thread_counts.py b/tests/python/test_thread_counts.py index c2b5b0c1..4896d8a9 100644 --- a/tests/python/test_thread_counts.py +++ b/tests/python/test_thread_counts.py @@ -1,5 +1,9 @@ """Thread requests are safe and preserve search results across all indexes.""" +import os +import subprocess +import sys + import numpy as np import pytest from rabitqlib import HnswIndex, IvfIndex, SymqgIndex @@ -52,3 +56,25 @@ def test_build_and_search_thread_counts(kind, metric, build_threads): ) np.testing.assert_array_equal(ids, expected_ids[:count]) np.testing.assert_array_equal(distances, expected_distances[:count]) + + +@pytest.mark.parametrize("limit", ["1", "2"]) +def test_automatic_thread_count_respects_omp_num_threads(limit): + # txtai sets OMP_NUM_THREADS=1 on macOS so coexisting OpenMP runtimes (Faiss, + # PyTorch) stay single-threaded; a default add must not start a worker team. + code = """ +import sys +from rabitqlib._rabitqlib import _available_cpu_count, _default_thread_count +expected = min(int(sys.argv[1]), _available_cpu_count()) +assert _default_thread_count() == expected, (_default_thread_count(), expected) +""" + env = os.environ.copy() + env["OMP_NUM_THREADS"] = limit + result = subprocess.run( + [sys.executable, "-c", code, limit], + env=env, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/unit/rabitqlib/utils/thread_count_test.cpp b/tests/unit/rabitqlib/utils/thread_count_test.cpp index dd8f81d3..27257ebe 100644 --- a/tests/unit/rabitqlib/utils/thread_count_test.cpp +++ b/tests/unit/rabitqlib/utils/thread_count_test.cpp @@ -19,7 +19,10 @@ TEST(ThreadCount, AutomaticAndOversizedRequestsUseAvailableCount) { const size_t available = rabitqlib::total_threads(); ASSERT_GE(available, 1U); - EXPECT_EQ(rabitqlib::resolve_num_threads(0), available); + EXPECT_EQ( + rabitqlib::resolve_num_threads(0), + std::min(available, static_cast(omp_get_max_threads())) + ); EXPECT_EQ(rabitqlib::resolve_num_threads(1), 1U); EXPECT_EQ(rabitqlib::resolve_num_threads(available), available); EXPECT_EQ(rabitqlib::resolve_num_threads(available + 1), available); @@ -31,6 +34,37 @@ TEST(ThreadCount, AutomaticAndOversizedRequestsUseAvailableCount) { } } +// Applications such as txtai set OMP_NUM_THREADS=1 on macOS so that coexisting +// OpenMP runtimes stay single-threaded; automatic teams must respect that limit. +TEST(ThreadCount, AutomaticRequestsRespectOpenMpThreadLimit) { + const int original = omp_get_max_threads(); + struct RestoreLimit { + int limit; + ~RestoreLimit() { omp_set_num_threads(limit); } + } restore{original}; + const size_t available = rabitqlib::total_threads(); + + omp_set_num_threads(1); + EXPECT_EQ(rabitqlib::default_threads(), 1U); + EXPECT_EQ(rabitqlib::resolve_num_threads(0), 1U); + EXPECT_EQ(rabitqlib::detail::openmp_team_size(0, 100), 1); + rabitqlib::index_detail::parallel_for(0, 19, 0, [](size_t, size_t worker) { + EXPECT_EQ(worker, 0U); + }); + // Explicit requests are capped only by the available CPUs. + EXPECT_EQ(rabitqlib::resolve_num_threads(2), std::min(2, available)); + EXPECT_EQ( + rabitqlib::resolve_num_threads(std::numeric_limits::max()), available + ); + + if (available >= 2) { + omp_set_num_threads(2); + EXPECT_EQ(rabitqlib::resolve_num_threads(0), 2U); + } + omp_set_num_threads(static_cast(available) + 1); + EXPECT_EQ(rabitqlib::resolve_num_threads(0), available); +} + #if defined(__linux__) TEST(ThreadCount, AutomaticAndExplicitRequestsRespectAffinity) { #if defined(_OPENMP) && _OPENMP >= 201307 @@ -58,8 +92,11 @@ TEST(ThreadCount, AutomaticAndExplicitRequestsRespectAffinity) { ASSERT_GT(available, 0U); ASSERT_EQ(sched_setaffinity(0, sizeof(restricted), &restricted), 0); EXPECT_EQ(rabitqlib::total_threads(), available); + const size_t automatic = + std::min(available, static_cast(omp_get_max_threads())); for (size_t requested : {size_t{0}, size_t{1}, std::numeric_limits::max()}) { - const size_t expected = requested == 1 ? 1 : available; + const size_t expected = + requested == 1 ? 1 : (requested == 0 ? automatic : available); EXPECT_EQ(rabitqlib::resolve_num_threads(requested), expected); rabitqlib::index_detail::parallel_for(0, 19, requested, [&](size_t, size_t worker) { EXPECT_LT(worker, expected); @@ -272,7 +309,13 @@ TEST(ThreadCount, OpenMpTeamSizeIsPositiveAndBounded) { EXPECT_EQ(rabitqlib::detail::openmp_team_size(1, 100), 1); EXPECT_EQ(rabitqlib::detail::openmp_team_size(0, 1), 1); EXPECT_EQ( - rabitqlib::detail::openmp_team_size(0, std::numeric_limits::max()), + rabitqlib::detail::openmp_team_size( + std::numeric_limits::max(), std::numeric_limits::max() + ), static_cast(std::min(available, std::numeric_limits::max())) ); + EXPECT_EQ( + rabitqlib::detail::openmp_team_size(0, std::numeric_limits::max()), + static_cast(rabitqlib::default_threads()) + ); }