diff --git a/THREAD_POOL_BEHAVIOR.md b/THREAD_POOL_BEHAVIOR.md new file mode 100644 index 000000000..e61471626 --- /dev/null +++ b/THREAD_POOL_BEHAVIOR.md @@ -0,0 +1,338 @@ +# ThreadPool Constructor and Configuration Behavior + +## Why this changed + +`ThreadPool` previously exposed both of these constructors: + +```cpp +ThreadPool(uint32_t size, bool binding); +ThreadPool(bool binding); +``` + +Calls such as `ThreadPool(1)` selected the exact `bool` overload instead of +converting `1` to `uint32_t`. As a result, a caller intending to create one +thread actually created `hardware_concurrency()` threads with CPU binding +enabled. Any positive one-argument thread count had the same problem. + +The ambiguous overload was removed. Thread count and CPU binding must now be +specified together: + +```cpp +ThreadPool(uint32_t size, bool binding); +``` + +## Direct C++ API + +| Construction | Thread count | CPU binding | +|---|---:|---| +| `ThreadPool()` | `max(std::thread::hardware_concurrency(), 1)` | Disabled | +| `ThreadPool(n, false)` | `n` | Disabled | +| `ThreadPool(n, true)` | `n` | Enabled | + +The default constructor has not changed: it uses hardware concurrency, never +zero, with binding disabled. + +The following one-argument forms are no longer supported: + +```cpp +ThreadPool(n); +ThreadPool(true); +ThreadPool(false); +``` + +Use explicit replacements: + +```cpp +// Old: ThreadPool(false) +ThreadPool(); + +// Old: ThreadPool(true) +ThreadPool(std::max(std::thread::hardware_concurrency(), 1u), true); + +// Intended meaning of old ThreadPool(n) +ThreadPool(n, false); // or true when binding is required +``` + +This is a C++ source compatibility change. It intentionally turns ambiguous +calls into compile errors instead of silently choosing the wrong behavior. + +CPU binding is implemented on Linux excluding Android. On unsupported +platforms, the binding flag does not change thread affinity. Linux binding +enumerates the calling thread's allowed CPU mask, so non-zero-based and +restricted cpusets are honored. Affinity API failures are logged. + +## DB global query and optimize pools + +The DB pools do not use the direct `ThreadPool()` default. Their counts are +initialized independently from `CgroupUtil::getCpuLimit()`: + +| Configuration | Default count | Default binding | +|---|---:|---| +| Query pool | Cgroup CPU limit | Enabled | +| Optimize pool | Cgroup CPU limit | Enabled | + +Therefore: + +- Omitting a DB thread count does **not** mean one thread. +- On an unrestricted host, the cgroup CPU limit will commonly equal available + CPU capacity. +- In a CPU-limited container, the DB defaults may differ from + `std::thread::hardware_concurrency()`. +- Cgroup v2 `cpu.max` values are parsed as `quota period`; `max period` is + treated as unlimited and falls back to the online CPU count. +- Query and optimize counts are independent. Setting `query_threads` does not + implicitly change `optimize_threads`. + +Before this fix, the DB passed its configured count as a single constructor +argument. Every valid positive count was converted to `true`, so the actual +runtime behavior was hardware-concurrency threads with binding enabled. This +meant configured values such as `optimize_threads=1` were ignored. + +After this fix, the configured count and binding are both forwarded +explicitly, so the effective behavior matches the configuration. Binding +remains enabled by default to preserve the previous effective behavior; users +can explicitly disable it when scheduler-managed placement is preferable. + +## Python API + +`zvec.init()` exposes the four independent options: + +```python +zvec.init( + query_threads=None, + query_thread_binding=None, + optimize_threads=None, + optimize_thread_binding=None, +) +``` + +Behavior: + +| Input | Effective behavior | +|---|---| +| Count is `None` | Use the corresponding cgroup-derived DB default | +| Count is a positive integer | Use exactly that many threads | +| Binding is `None` | Use the DB default (`True`) | +| Binding is `False` | Do not bind worker threads | +| Binding is `True` | Bind worker threads where supported | + +Examples: + +```python +# Both pools use their cgroup-derived counts with binding. +zvec.init() + +# Optimize uses exactly one bound worker. +zvec.init(optimize_threads=1) + +# Explicitly let the OS schedule both pools without affinity pinning. +zvec.init(query_thread_binding=False, optimize_thread_binding=False) +``` + +## C API + +The C configuration API now includes independent setters and getters: + +```c +zvec_config_data_set_query_thread_binding(config, binding); +zvec_config_data_get_query_thread_binding(config); +zvec_config_data_set_optimize_thread_binding(config, binding); +zvec_config_data_get_optimize_thread_binding(config); +``` + +If these setters are not called, binding remains enabled. + +## Internal call-site migration + +Internal call sites were migrated according to their previous intended +behavior: + +- Benchmark tools with no requested count still use hardware concurrency + without binding. +- Recall tools with no requested count still use hardware concurrency with + binding. +- Explicit benchmark and recall thread counts are now honored. +- `Index::Merge` uses `write_concurrency` with binding, preserving its previous + effective affinity behavior. +- DB query and optimize pools use their configured count and binding. +- Existing call sites already passing both count and binding are unchanged. + +The recall resize path previously used the arguments in reverse order: + +```cpp +ThreadPool(true, threads); +``` + +That constructed one bound worker. It now correctly uses: + +```cpp +ThreadPool(threads, true); +``` + +## Performance validation + +The fix was validated with an end-to-end Python DB API benchmark using the +SIFT dataset. The measured post-fix revision used disabled binding defaults. +The current compatibility-preserving defaults keep binding enabled, so the +post-fix placement measured below is reproduced by explicitly setting both +binding options to `False`. + +| Parameter | Value | +|---|---| +| Dataset | SIFT, 1,000,000 FP32 vectors, 128 dimensions | +| Index | Vamana | +| Search list size | 500 | +| Alpha | 1.5 | +| Quantizer | Uniform INT8 | +| Optimize threads | 1 per process | +| Build flow | `create_and_open -> insert -> optimize` | + +Two codebases were compared on the same machine with separately built wheels +and virtual environments: + +- Before the fix: commit `a90ec5b1d658d`, wheel `0.5.2.dev30`. +- After the fix: the ThreadPool fix and migrated call sites, wheel + `0.5.2.dev32`. + +### Observed runtime behavior + +The test machine exposes 64 logical CPUs. The benchmark called: + +```python +zvec.init(optimize_threads=1) +``` + +It did not override the query thread count or either binding option. In that +measured revision the effective binding defaults were `False`; with the +current defaults, use the following equivalent configuration for the post-fix +run: + +```python +zvec.init( + optimize_threads=1, + query_thread_binding=False, + optimize_thread_binding=False, +) +``` + +| Pool / behavior | Before the fix | After the fix | +|---|---|---| +| Query pool | 64 workers, bound to CPUs 0-63 | 64 workers, unbound | +| Optimize pool | 64 workers, bound to CPUs 0-63 | 1 worker, unbound | +| Total pool workers | 128 | 65 | +| Observed process threads | Approximately 214 | Approximately 150 | +| Workers doing Vamana build work | Effectively about 1 per process | 1 per process | +| CPU used by each builder in the four-process phase | Approximately 25% of one CPU before competing builders completed | Approximately 92-100% of one CPU | + +The total process thread count is higher than the two global pools because it +also includes Python, storage-engine, I/O, and other background threads. Those +additional threads were present in both versions. + +Before the fix, both positive configuration values were passed through the +one-argument constructor: + +```cpp +ThreadPool(configured_count); +``` + +Both `64` and `1` converted to the Boolean value `true`. Consequently, both +pools created 64 workers and enabled binding. `optimize_threads=1` therefore +did not create one optimize worker. + +The Vamana build remained effectively single-threaded in this test: despite +the 64 optimize workers, only about one worker per process was doing sustained +CPU work. The other optimize workers and the entire query pool were mostly +idle. This explains why the single-process result did not become faster in the +old version. + +The excessive idle workers add thread stacks, scheduler bookkeeping, and some +context-switch overhead, but they are not sufficient to explain the observed +four-times slowdown. The dominant problem was CPU affinity: + +1. Every process created the same 64-worker optimize pool. +2. Worker `i` in every process was bound to CPU `i`. +3. The effectively active worker selected by each identical pool tended to + have the same worker index. +4. Active build work from multiple processes therefore competed on the same + CPU instead of being spread across the 64 available CPUs. + +This behavior matches the timing curve. The degree-24 build, which overlapped +with all four processes for almost its entire lifetime, took 4.28 times its +single-process time before the fix. As shorter jobs completed, the remaining +processes received a larger share of the contended CPU, so the longest +degree-64 job showed a smaller but still substantial 2.45-times slowdown. + +After the fix, each process created one unbound optimize worker. The operating +system could schedule the four runnable workers on different CPUs, and live +sampling showed each builder consuming close to one full CPU. The query pool +still contributed 64 mostly idle workers because the benchmark did not set +`query_threads`; this accounts for much of the approximately 150-thread process +total, but did not prevent scaling because those threads were unbound and idle. + +### Single-process results + +| Degree | Before | After | Change | +|---:|---:|---:|---:| +| 24 | 790.2 s | 777.8 s | -1.6% | +| 64 | 2865.3 s | 2847.0 s | -0.6% | + +Single-process performance was effectively unchanged. This indicates that the +fix does not alter Vamana's underlying build cost or introduce a measurable +single-process regression. + +### Four-process results + +Four independent processes simultaneously built indexes with degrees 24, 32, +48, and 64. Every process requested `optimize_threads=1`. + +| Degree | Before | After | Speedup | Time reduction | +|---:|---:|---:|---:|---:| +| 24 | 3378.6 s | 841.6 s | 4.01x | 75.1% | +| 32 | 4282.3 s | 1137.5 s | 3.76x | 73.4% | +| 48 | 5955.2 s | 1946.4 s | 3.06x | 67.3% | +| 64 | 7009.0 s | 2975.1 s | 2.36x | 57.6% | + +The four-process group wall-clock time was determined by the degree-64 build: + +| Metric | Before | After | Improvement | +|---|---:|---:|---:| +| Group wall time | 7009.2 s (116m 49s) | 2975.4 s (49m 35s) | 57.5% lower | +| Equivalent speedup | — | — | 2.36x | + +### Parallel overhead relative to a single process + +| Degree | Before fix | After fix | +|---:|---:|---:| +| 24 | +327.6% | +8.2% | +| 64 | +144.6% | +4.5% | + +Before the fix, the degree-24 build took 4.28 times as long when sharing the +machine with three other builders. After the fix it took only 1.08 times as +long. The degree-64 ratio similarly improved from 2.45 times to 1.05 times. + +All indexes produced before and after the fix reached an index completeness of +1.0. The results confirm that the change specifically removes multi-process +thread contention while preserving single-process performance and index +correctness. + +## Compatibility summary + +Behavior is unchanged for: + +- `ThreadPool()`; +- `ThreadPool(n, false)`; +- `ThreadPool(n, true)`; +- `Index::Merge` CPU binding; +- DB query and optimize pool binding defaults; +- internal call sites whose old behavior was already explicit and correct. + +Behavior intentionally changes for: + +- one-argument thread counts that were previously interpreted as booleans; +- DB thread-count configuration that was previously ignored; +- recall resize paths affected by reversed arguments. + +DB binding remains enabled by default, matching the previous effective +runtime behavior. Adding binding fields to the public C++ `ConfigData` struct +changes its layout, so consumers linking across a C++ ABI boundary must rebuild +against the updated headers. diff --git a/python/tests/detail/test_db_config.py b/python/tests/detail/test_db_config.py index 16203a8fa..e4f993339 100644 --- a/python/tests/detail/test_db_config.py +++ b/python/tests/detail/test_db_config.py @@ -153,6 +153,21 @@ class TestDbConfigThreadValidation: def test_query_threads(self): zvec.init(query_threads=1) + @run_in_subprocess + def test_query_thread_binding(self): + zvec.init(query_thread_binding=True) + + @run_in_subprocess + def test_query_thread_binding_disabled(self): + zvec.init(query_thread_binding=False) + + @run_in_subprocess + def test_query_thread_binding_invalid(self): + with pytest.raises(TypeError): + zvec.init(query_thread_binding=1) + with pytest.raises(TypeError): + zvec.init(query_thread_binding="true") + @run_in_subprocess def test_query_threads_invalid(self): # query_threads must >= 0 and must be int and if None, set default value @@ -171,6 +186,21 @@ def test_query_threads_invalid(self): def test_optimize_threads(self): zvec.init(optimize_threads=1) + @run_in_subprocess + def test_optimize_thread_binding(self): + zvec.init(optimize_thread_binding=True) + + @run_in_subprocess + def test_optimize_thread_binding_disabled(self): + zvec.init(optimize_thread_binding=False) + + @run_in_subprocess + def test_optimize_thread_binding_invalid(self): + with pytest.raises(TypeError): + zvec.init(optimize_thread_binding=1) + with pytest.raises(TypeError): + zvec.init(optimize_thread_binding="true") + @run_in_subprocess def test_optimize_threads_invalid(self): # optimize_threads must >= 0 and must be int and if None, set default value diff --git a/python/zvec/zvec.py b/python/zvec/zvec.py index 0046a496d..f2953236c 100644 --- a/python/zvec/zvec.py +++ b/python/zvec/zvec.py @@ -35,7 +35,9 @@ def init( log_file_size: Optional[int] = 2048, log_overdue_days: Optional[int] = 7, query_threads: Optional[int] = None, + query_thread_binding: Optional[bool] = None, optimize_threads: Optional[int] = None, + optimize_thread_binding: Optional[bool] = None, invert_to_forward_scan_ratio: Optional[float] = None, brute_force_by_keys_ratio: Optional[float] = None, fts_brute_force_by_keys_ratio: Optional[float] = None, @@ -79,9 +81,16 @@ def init( Number of threads for query execution. If ``None`` (default), inferred from available CPU cores (via cgroup). Must be ≥ 1 if provided. + query_thread_binding (Optional[bool], optional): + Whether to bind query worker threads to CPU cores. Defaults to + Zvec's internal setting, currently ``True``. optimize_threads (Optional[int], optional): Threads for background tasks (e.g., compaction, indexing). - If ``None``, defaults to same as ``query_threads`` or CPU count. + If ``None`` (default), inferred independently from available CPU + cores (via cgroup). + optimize_thread_binding (Optional[bool], optional): + Whether to bind optimize worker threads to CPU cores. Defaults to + Zvec's internal setting, currently ``True``. invert_to_forward_scan_ratio (Optional[float], optional): Threshold to switch from inverted index to full forward scan. Range: [0.0, 1.0]. Higher → more aggressive index skipping. @@ -167,8 +176,16 @@ def init( config_dict["log_overdue_days"] = log_overdue_days if query_threads is not None: config_dict["query_threads"] = query_threads + if query_thread_binding is not None: + if not isinstance(query_thread_binding, bool): + raise TypeError("query_thread_binding must be bool") + config_dict["query_thread_binding"] = query_thread_binding if optimize_threads is not None: config_dict["optimize_threads"] = optimize_threads + if optimize_thread_binding is not None: + if not isinstance(optimize_thread_binding, bool): + raise TypeError("optimize_thread_binding must be bool") + config_dict["optimize_thread_binding"] = optimize_thread_binding if invert_to_forward_scan_ratio is not None: config_dict["invert_to_forward_scan_ratio"] = invert_to_forward_scan_ratio if brute_force_by_keys_ratio is not None: diff --git a/src/ailego/parallel/thread_pool.cc b/src/ailego/parallel/thread_pool.cc index 54a8970e4..656b99ead 100644 --- a/src/ailego/parallel/thread_pool.cc +++ b/src/ailego/parallel/thread_pool.cc @@ -12,20 +12,47 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include -#if (defined(__linux) || defined(__linux__)) && !defined(__ANDROID__) +#if defined(__linux__) && !defined(__ANDROID__) #include +static inline bool GetAllowedCpuMask(cpu_set_t *mask) { + CPU_ZERO(mask); + const int error = pthread_getaffinity_np(pthread_self(), sizeof(*mask), mask); + if (error != 0) { + LOG_WARN("Failed to get thread affinity mask, error[%d]", error); + return false; + } + return true; +} + static inline void BindThreads(std::vector &pool) { - uint32_t hc = std::thread::hardware_concurrency(); - if (hc > 1) { - cpu_set_t mask; - - for (size_t i = 0u; i < pool.size(); ++i) { - CPU_ZERO(&mask); - CPU_SET(i % hc, &mask); - pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask); + cpu_set_t allowed_mask; + if (!GetAllowedCpuMask(&allowed_mask)) { + return; + } + + std::vector allowed_cpus; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &allowed_mask)) { + allowed_cpus.push_back(cpu); + } + } + if (allowed_cpus.empty()) { + LOG_WARN("Cannot bind thread pool: affinity mask has no allowed CPUs"); + return; + } + + cpu_set_t target_mask; + for (size_t i = 0u; i < pool.size(); ++i) { + CPU_ZERO(&target_mask); + CPU_SET(allowed_cpus[i % allowed_cpus.size()], &target_mask); + const int error = pthread_setaffinity_np(pool[i].native_handle(), + sizeof(target_mask), &target_mask); + if (error != 0) { + LOG_WARN("Failed to bind thread pool worker[%zu], error[%d]", i, error); } } } @@ -33,12 +60,16 @@ static inline void BindThreads(std::vector &pool) { static inline void UnbindThreads(std::vector &pool) { cpu_set_t mask; CPU_ZERO(&mask); - - for (size_t i = 0u; i < CPU_SETSIZE; ++i) { - CPU_SET(i, &mask); + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + CPU_SET(cpu, &mask); } + for (size_t i = 0u; i < pool.size(); ++i) { - pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask); + const int error = + pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask); + if (error != 0) { + LOG_WARN("Failed to unbind thread pool worker[%zu], error[%d]", i, error); + } } } #else @@ -130,4 +161,4 @@ bool ThreadPool::picking(ThreadPool::Task *task) { } } // namespace ailego -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/binding/c/c_api.cc b/src/binding/c/c_api.cc index 1dfce56f2..1a86dc5cf 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -589,6 +589,27 @@ uint32_t zvec_config_data_get_query_thread_count(const zvec_config_data_t *confi return cpp_config->query_thread_count; } +zvec_error_code_t zvec_config_data_set_query_thread_binding( + zvec_config_data_t *config, bool binding) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->query_thread_binding = binding; + return ZVEC_OK; +} + +bool zvec_config_data_get_query_thread_binding( + const zvec_config_data_t *config) { + if (!config) { + return false; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->query_thread_binding; +} + zvec_error_code_t zvec_config_data_set_invert_to_forward_scan_ratio( zvec_config_data_t *config, float ratio) { if (!config) { @@ -673,6 +694,27 @@ uint32_t zvec_config_data_get_optimize_thread_count( return cpp_config->optimize_thread_count; } +zvec_error_code_t zvec_config_data_set_optimize_thread_binding( + zvec_config_data_t *config, bool binding) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->optimize_thread_binding = binding; + return ZVEC_OK; +} + +bool zvec_config_data_get_optimize_thread_binding( + const zvec_config_data_t *config) { + if (!config) { + return false; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->optimize_thread_binding; +} + zvec_error_code_t zvec_config_data_set_jieba_dict_dir( zvec_config_data_t *config, const char *dir) { if (!config) { diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index d6ad1f48b..a8e75046b 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -150,6 +150,12 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { data.query_thread_count = static_cast(q); } + // set query thread binding + if (has_key(config_dict, "query_thread_binding")) { + data.query_thread_binding = + get_if(config_dict, "query_thread_binding").value(); + } + // set optimize thread count if (has_key(config_dict, "optimize_threads")) { auto o = get_if(config_dict, "optimize_threads").value(); @@ -157,6 +163,12 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { data.optimize_thread_count = static_cast(o); } + // set optimize thread binding + if (has_key(config_dict, "optimize_thread_binding")) { + data.optimize_thread_binding = + get_if(config_dict, "optimize_thread_binding").value(); + } + // set invert_to_forward_scan_ratio if (has_key(config_dict, "invert_to_forward_scan_ratio")) { auto v = diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 755f673f9..2187a8bad 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -986,7 +986,7 @@ int Index::Merge(const std::vector &indexes, reducer->set_thread_pool(options.pool); } else { local_thread_pool = - std::make_unique(options.write_concurrency); + std::make_unique(options.write_concurrency, true); reducer->set_thread_pool(local_thread_pool.get()); } diff --git a/src/db/common/cgroup_util.cc b/src/db/common/cgroup_util.cc index 0e137d75d..a9e4b905d 100644 --- a/src/db/common/cgroup_util.cc +++ b/src/db/common/cgroup_util.cc @@ -13,6 +13,8 @@ // limitations under the License. #include "db/common/cgroup_util.h" +#include +#include namespace zvec { @@ -114,19 +116,15 @@ bool CgroupUtil::readCpuCgroup() { // cgroup v2 std::ifstream file("/sys/fs/cgroup/cpu.max"); if (file.is_open()) { - uint64_t quota, period; - char slash; - file >> quota >> slash >> period; - file.close(); + std::string cpu_max; + std::getline(file, cpu_max); - if (quota != std::numeric_limits::max() && quota != 0 && - period > 0) { - cpu_cores_ = - static_cast(std::ceil(static_cast(quota) / period)); + int cpu_cores = 0; + if (parseCpuMax(cpu_max, &cpu_cores)) { + cpu_cores_ = cpu_cores; return true; - } else { - return false; } + return false; } // cgroup v1 @@ -150,6 +148,43 @@ bool CgroupUtil::readCpuCgroup() { return false; } +bool CgroupUtil::parseCpuMax(const std::string &cpu_max, int *cpu_cores) { + if (cpu_cores == nullptr) { + return false; + } + + std::istringstream stream(cpu_max); + std::string quota_token; + uint64_t period = 0; + if (!(stream >> quota_token >> period) || period == 0) { + return false; + } + + std::string trailing_token; + if (stream >> trailing_token) { + return false; + } + if (quota_token == "max") { + return false; + } + + uint64_t quota = 0; + const auto result = std::from_chars( + quota_token.data(), quota_token.data() + quota_token.size(), quota); + if (result.ec != std::errc{} || + result.ptr != quota_token.data() + quota_token.size() || quota == 0) { + return false; + } + + const uint64_t cores = quota / period + (quota % period != 0 ? 1 : 0); + if (cores > static_cast(std::numeric_limits::max())) { + return false; + } + + *cpu_cores = static_cast(cores); + return true; +} + void CgroupUtil::updateMemoryLimit() { if (readMemoryCgroup()) { return; @@ -491,4 +526,4 @@ double CgroupUtil::calculateMacOSCpuUsage() { } #endif -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/common/cgroup_util.h b/src/db/common/cgroup_util.h index 759e2317c..f7d1d70ad 100644 --- a/src/db/common/cgroup_util.h +++ b/src/db/common/cgroup_util.h @@ -45,6 +45,10 @@ class CgroupUtil { static int getCpuLimit(); static uint64_t getMemoryLimit(); + // Parse a cgroup v2 cpu.max value. Returns false for an unlimited or + // malformed value so callers can fall back to the host CPU count. + static bool parseCpuMax(const std::string &cpu_max, int *cpu_cores); + // Static methods to get other resources static double getCpuUsage(); static uint64_t getMemoryUsage(); @@ -101,4 +105,4 @@ class CgroupUtil { #endif }; -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/common/config.cc b/src/db/common/config.cc index 57eaae812..0915294b7 100644 --- a/src/db/common/config.cc +++ b/src/db/common/config.cc @@ -35,10 +35,12 @@ GlobalConfig::ConfigData::ConfigData() DEFAULT_MEMORY_LIMIT_RATIO), log_config(std::make_shared()), query_thread_count(CgroupUtil::getCpuLimit()), + query_thread_binding(true), invert_to_forward_scan_ratio(0.9), brute_force_by_keys_ratio(0.1), fts_brute_force_by_keys_ratio(0.05), optimize_thread_count(CgroupUtil::getCpuLimit()), + optimize_thread_binding(true), jieba_dict_dir() {} Status GlobalConfig::Validate(const ConfigData &config) const { @@ -165,4 +167,4 @@ uint64_t GlobalConfig::memory_limit_bytes() const noexcept { FACTORY_REGISTER_LOGGER(AppendLogger); -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/common/global_resource.cc b/src/db/common/global_resource.cc index 711bf3dbd..f15bce8d5 100644 --- a/src/db/common/global_resource.cc +++ b/src/db/common/global_resource.cc @@ -21,10 +21,12 @@ namespace zvec { void GlobalResource::initialize() { static std::once_flag flag; std::call_once(flag, [this]() mutable { - this->query_thread_pool_.reset( - new ailego::ThreadPool(GlobalConfig::Instance().query_thread_count())); + this->query_thread_pool_.reset(new ailego::ThreadPool( + GlobalConfig::Instance().query_thread_count(), + GlobalConfig::Instance().query_thread_binding())); this->optimize_thread_pool_.reset(new ailego::ThreadPool( - GlobalConfig::Instance().optimize_thread_count())); + GlobalConfig::Instance().optimize_thread_count(), + GlobalConfig::Instance().optimize_thread_binding())); zvec::ailego::MemoryLimitPool::get_instance().init( GlobalConfig::Instance().memory_limit_bytes()); }); diff --git a/src/include/zvec/ailego/parallel/thread_pool.h b/src/include/zvec/ailego/parallel/thread_pool.h index 150f014c0..36542dae4 100644 --- a/src/include/zvec/ailego/parallel/thread_pool.h +++ b/src/include/zvec/ailego/parallel/thread_pool.h @@ -137,12 +137,8 @@ class ThreadPool { explicit ThreadPool(uint32_t size, bool binding); //! Constructor - explicit ThreadPool(bool binding) - : ThreadPool{std::max(std::thread::hardware_concurrency(), 1u), binding} { - } - - //! Constructor - ThreadPool(void) : ThreadPool{false} {} + ThreadPool(void) + : ThreadPool{std::max(std::thread::hardware_concurrency(), 1u), false} {} //! Destructor ~ThreadPool(void) { diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 2d048689b..ca496d8d3 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -643,6 +643,25 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_config_data_set_query_thread_count( ZVEC_EXPORT uint32_t ZVEC_CALL zvec_config_data_get_query_thread_count(const zvec_config_data_t *config); +/** + * @brief Set query thread CPU binding in configuration data + * @param config Configuration data pointer + * @param binding Whether query threads should bind to CPU cores + * The default is true when this setter is not called. + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_config_data_set_query_thread_binding(zvec_config_data_t *config, + bool binding); + +/** + * @brief Get query thread CPU binding from configuration data + * @param config Configuration data pointer + * @return bool Query thread CPU binding + */ +ZVEC_EXPORT bool ZVEC_CALL +zvec_config_data_get_query_thread_binding(const zvec_config_data_t *config); + /** * @brief Set invert to forward scan ratio in configuration data * @param config Configuration data pointer @@ -715,6 +734,25 @@ zvec_config_data_set_optimize_thread_count(zvec_config_data_t *config, ZVEC_EXPORT uint32_t ZVEC_CALL zvec_config_data_get_optimize_thread_count(const zvec_config_data_t *config); +/** + * @brief Set optimize thread CPU binding in configuration data + * @param config Configuration data pointer + * @param binding Whether optimize threads should bind to CPU cores + * The default is true when this setter is not called. + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_config_data_set_optimize_thread_binding(zvec_config_data_t *config, + bool binding); + +/** + * @brief Get optimize thread CPU binding from configuration data + * @param config Configuration data pointer + * @return bool Optimize thread CPU binding + */ +ZVEC_EXPORT bool ZVEC_CALL +zvec_config_data_get_optimize_thread_binding(const zvec_config_data_t *config); + /** * @brief Set jieba dict directory in configuration data * @param dir Dict directory; NULL or empty leaves the field empty diff --git a/src/include/zvec/db/config.h b/src/include/zvec/db/config.h index 4403f352b..07836b98b 100644 --- a/src/include/zvec/db/config.h +++ b/src/include/zvec/db/config.h @@ -92,6 +92,8 @@ class GlobalConfig : public ailego::Singleton { // query uint32_t query_thread_count; + // Preserve the historical effective behavior unless explicitly disabled. + bool query_thread_binding; float invert_to_forward_scan_ratio; float brute_force_by_keys_ratio; // Independent from brute_force_by_keys_ratio: per-candidate FTS cost @@ -100,6 +102,8 @@ class GlobalConfig : public ailego::Singleton { // optimize uint32_t optimize_thread_count; + // Preserve the historical effective behavior unless explicitly disabled. + bool optimize_thread_binding; // FTS jieba tokenizer default dict dir (lowest-priority fallback; // per-field config > ZVEC_JIEBA_DICT_DIR > this). Empty by default. @@ -165,6 +169,11 @@ class GlobalConfig : public ailego::Singleton { return config_.query_thread_count; } + //! Query thread binding + bool query_thread_binding() const noexcept { + return config_.query_thread_binding; + } + //! Invert to forward scan ratio float invert_to_forward_scan_ratio() const noexcept { return config_.invert_to_forward_scan_ratio; @@ -186,6 +195,11 @@ class GlobalConfig : public ailego::Singleton { return config_.optimize_thread_count; } + //! Optimize thread binding + bool optimize_thread_binding() const noexcept { + return config_.optimize_thread_binding; + } + //! Effective jieba dict dir. Thread-safe. std::string jieba_dict_dir() const; diff --git a/tests/ailego/parallel/thread_pool_test.cc b/tests/ailego/parallel/thread_pool_test.cc index 70459f535..23226cd15 100644 --- a/tests/ailego/parallel/thread_pool_test.cc +++ b/tests/ailego/parallel/thread_pool_test.cc @@ -15,11 +15,79 @@ #include #include #include +#include #include #include +#if defined(__linux__) && !defined(__ANDROID__) +#include +#endif + using namespace zvec::ailego; +static_assert(!std::is_constructible::value, + "ThreadPool thread count must not be ambiguous with binding"); +static_assert(!std::is_constructible::value, + "ThreadPool binding must require an explicit thread count"); +static_assert(std::is_constructible::value, + "ThreadPool must accept an explicit count and binding"); + +TEST(ThreadPool, ConstructorBehavior) { + const auto hardware_concurrency = + std::max(std::thread::hardware_concurrency(), 1u); + + ThreadPool default_pool; + EXPECT_EQ(hardware_concurrency, default_pool.count()); + + ThreadPool single_unbound_pool(1, false); + EXPECT_EQ(1u, single_unbound_pool.count()); + + ThreadPool double_bound_pool(2, true); + EXPECT_EQ(2u, double_bound_pool.count()); +} + +#if defined(__linux__) && !defined(__ANDROID__) +TEST(ThreadPool, BindingRespectsCallerAffinityMask) { + cpu_set_t caller_mask; + CPU_ZERO(&caller_mask); + ASSERT_EQ(0, pthread_getaffinity_np(pthread_self(), sizeof(caller_mask), + &caller_mask)); + ASSERT_GT(CPU_COUNT(&caller_mask), 0); + + ThreadPool pool(1, true); + cpu_set_t bound_mask; + CPU_ZERO(&bound_mask); + int get_affinity_error = -1; + pool.execute_and_wait([&]() { + get_affinity_error = + pthread_getaffinity_np(pthread_self(), sizeof(bound_mask), &bound_mask); + }); + + ASSERT_EQ(0, get_affinity_error); + ASSERT_EQ(1, CPU_COUNT(&bound_mask)); + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &bound_mask)) { + EXPECT_TRUE(CPU_ISSET(cpu, &caller_mask)); + } + } + + pool.unbind(); + cpu_set_t unbound_mask; + CPU_ZERO(&unbound_mask); + pool.execute_and_wait([&]() { + get_affinity_error = pthread_getaffinity_np( + pthread_self(), sizeof(unbound_mask), &unbound_mask); + }); + + ASSERT_EQ(0, get_affinity_error); + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &caller_mask)) { + EXPECT_TRUE(CPU_ISSET(cpu, &unbound_mask)); + } + } +} +#endif + struct A { A(void) : pool(std::make_shared()) {} @@ -38,7 +106,9 @@ struct A { }; struct B { - B(void) : pool(std::make_shared(true)) {} + B(void) + : pool(std::make_shared( + std::max(std::thread::hardware_concurrency(), 1u), true)) {} std::string ThreadMain(uint32_t &num) { aaa.pool->enqueue( diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 7365dcf1a..4f6eb9e67 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -234,10 +234,20 @@ void test_zvec_config() { TEST_ASSERT(err == ZVEC_OK); TEST_ASSERT(zvec_config_data_get_query_thread_count(config_data) == 8); + TEST_ASSERT(zvec_config_data_get_query_thread_binding(config_data)); + err = zvec_config_data_set_query_thread_binding(config_data, false); + TEST_ASSERT(err == ZVEC_OK); + TEST_ASSERT(!zvec_config_data_get_query_thread_binding(config_data)); + err = zvec_config_data_set_optimize_thread_count(config_data, 4); TEST_ASSERT(err == ZVEC_OK); TEST_ASSERT(zvec_config_data_get_optimize_thread_count(config_data) == 4); + TEST_ASSERT(zvec_config_data_get_optimize_thread_binding(config_data)); + err = zvec_config_data_set_optimize_thread_binding(config_data, false); + TEST_ASSERT(err == ZVEC_OK); + TEST_ASSERT(!zvec_config_data_get_optimize_thread_binding(config_data)); + // Test log config replacement TEST_ASSERT(zvec_config_data_get_log_type(config_data) == ZVEC_LOG_TYPE_CONSOLE); @@ -263,9 +273,15 @@ void test_zvec_config() { err = zvec_config_data_set_query_thread_count(NULL, 1); TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + err = zvec_config_data_set_query_thread_binding(NULL, true); + TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + err = zvec_config_data_set_optimize_thread_count(NULL, 1); TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + err = zvec_config_data_set_optimize_thread_binding(NULL, true); + TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + // Test boundary values zvec_config_data_t *boundary_config = zvec_config_data_create(); if (boundary_config) { diff --git a/tests/db/common/cgroup_util_test.cc b/tests/db/common/cgroup_util_test.cc new file mode 100644 index 000000000..e011f8c6b --- /dev/null +++ b/tests/db/common/cgroup_util_test.cc @@ -0,0 +1,45 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "db/common/cgroup_util.h" +#include + +using namespace zvec; + +TEST(CgroupUtil, ParseCpuMaxQuota) { + int cpu_cores = 0; + + ASSERT_TRUE(CgroupUtil::parseCpuMax("100000 100000", &cpu_cores)); + EXPECT_EQ(1, cpu_cores); + + ASSERT_TRUE(CgroupUtil::parseCpuMax("150000 100000\n", &cpu_cores)); + EXPECT_EQ(2, cpu_cores); + + ASSERT_TRUE(CgroupUtil::parseCpuMax("400000 100000", &cpu_cores)); + EXPECT_EQ(4, cpu_cores); +} + +TEST(CgroupUtil, ParseCpuMaxUnlimitedOrInvalid) { + int cpu_cores = 7; + + EXPECT_FALSE(CgroupUtil::parseCpuMax("max 100000", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000/100000", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000 0", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("0 100000", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000 100000 trailing", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("invalid", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000 100000", nullptr)); + + EXPECT_EQ(7, cpu_cores); +} diff --git a/tests/db/common/config_test.cc b/tests/db/common/config_test.cc index d86e160d4..022664059 100644 --- a/tests/db/common/config_test.cc +++ b/tests/db/common/config_test.cc @@ -29,6 +29,15 @@ class ConfigTest : public ::testing::Test { } }; +TEST_F(ConfigTest, ThreadConfigDataDefaults) { + GlobalConfig::ConfigData config; + + ASSERT_GT(config.query_thread_count, 0u); + ASSERT_EQ(config.query_thread_count, config.optimize_thread_count); + ASSERT_TRUE(config.query_thread_binding); + ASSERT_TRUE(config.optimize_thread_binding); +} + TEST_F(ConfigTest, InitializeWithDefaultConfig) { GlobalConfig::ConfigData config; @@ -42,10 +51,12 @@ TEST_F(ConfigTest, InitializeWithDefaultConfig) { GlobalConfig::LogLevel::kWarn); ASSERT_EQ(GlobalConfig::Instance().log_type(), "ConsoleLogger"); ASSERT_GT(GlobalConfig::Instance().query_thread_count(), 0); + ASSERT_TRUE(GlobalConfig::Instance().query_thread_binding()); ASSERT_EQ(GlobalConfig::Instance().invert_to_forward_scan_ratio(), 0.9f); ASSERT_EQ(GlobalConfig::Instance().brute_force_by_keys_ratio(), 0.1f); ASSERT_EQ(GlobalConfig::Instance().fts_brute_force_by_keys_ratio(), 0.05f); ASSERT_GT(GlobalConfig::Instance().optimize_thread_count(), 0); + ASSERT_TRUE(GlobalConfig::Instance().optimize_thread_binding()); } TEST_F(ConfigTest, InitializeWithCustomConsoleLogConfig) { @@ -243,4 +254,4 @@ TEST_F(ConfigTest, JiebaDictDirSetterIsIndependentOfInitialize) { ASSERT_EQ(GlobalConfig::Instance().jieba_dict_dir(), ""); GlobalConfig::Instance().set_default_jieba_dict_dir(saved); -} \ No newline at end of file +} diff --git a/tests/db/common/global_resource_test.cc b/tests/db/common/global_resource_test.cc new file mode 100644 index 000000000..2f37d81f9 --- /dev/null +++ b/tests/db/common/global_resource_test.cc @@ -0,0 +1,33 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "db/common/global_resource.h" +#include +#include + +using namespace zvec; + +TEST(GlobalResource, UsesConfiguredCountsAndDefaultBinding) { + GlobalConfig::ConfigData config; + config.query_thread_count = 2; + config.optimize_thread_count = 1; + + const auto status = GlobalConfig::Instance().Initialize(config); + ASSERT_TRUE(status.ok()) << status.message(); + + EXPECT_EQ(2u, GlobalResource::Instance().query_thread_pool()->count()); + EXPECT_EQ(1u, GlobalResource::Instance().optimize_thread_pool()->count()); + EXPECT_TRUE(GlobalConfig::Instance().query_thread_binding()); + EXPECT_TRUE(GlobalConfig::Instance().optimize_thread_binding()); +} diff --git a/tools/core/bench.cc b/tools/core/bench.cc index 2587dfe2f..f7407e74a 100644 --- a/tools/core/bench.cc +++ b/tools/core/bench.cc @@ -35,7 +35,7 @@ class Bench { retrieval_mode_{retrieval_mode}, filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; @@ -339,7 +339,7 @@ class SparseBench { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; diff --git a/tools/core/bench_original.cc b/tools/core/bench_original.cc index c787b883b..f3b408862 100644 --- a/tools/core/bench_original.cc +++ b/tools/core/bench_original.cc @@ -69,7 +69,7 @@ class Bench { retrieval_mode_{retrieval_mode}, filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; @@ -380,7 +380,7 @@ class SparseBench { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; diff --git a/tools/core/recall.cc b/tools/core/recall.cc index d5be1919c..cace81bbc 100644 --- a/tools/core/recall.cc +++ b/tools/core/recall.cc @@ -36,7 +36,8 @@ class Recall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared( + std::max(std::thread::hardware_concurrency(), 1u), true); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; @@ -89,7 +90,7 @@ class Recall { if (batch_queries_.size() < threads_) { threads_ = batch_queries_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, true); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; } @@ -738,7 +739,8 @@ class SparseRecall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared( + std::max(std::thread::hardware_concurrency(), 1u), true); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; @@ -820,7 +822,7 @@ class SparseRecall { if (batch_sparse_counts_.size() < threads_) { threads_ = batch_sparse_counts_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, true); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; } diff --git a/tools/core/recall_original.cc b/tools/core/recall_original.cc index 997ef5ad7..2a0ce962d 100644 --- a/tools/core/recall_original.cc +++ b/tools/core/recall_original.cc @@ -76,7 +76,8 @@ class Recall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared( + std::max(std::thread::hardware_concurrency(), 1u), true); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; @@ -127,7 +128,7 @@ class Recall { if (batch_queries_.size() < threads_) { threads_ = batch_queries_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, true); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; } @@ -903,7 +904,8 @@ class SparseRecall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared( + std::max(std::thread::hardware_concurrency(), 1u), true); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; @@ -985,7 +987,7 @@ class SparseRecall { if (batch_sparse_counts_.size() < threads_) { threads_ = batch_sparse_counts_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, true); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; }