GH-45860: [C++] Respect CPU affinity in cpu_count and ThreadPool default capacity #46034

MuteBardTison · 2025-04-04T19:02:59Z

Rationale for this change

Arrow’s current CPU thread count detection uses std::thread::hardware_concurrency() which does not take into account the process-level CPU affinity mask (e.g., set via taskset). This can lead to thread oversubscription and performance issues when Arrow runs in constrained environments.

This PR improves Arrow's behavior on Linux by making both CpuInfo::num_cores() and ThreadPool::DefaultCapacity() aware of CPU affinity.

What changes are included in this PR?

Added affinity.h to expose GetAffinityCpuCount() on Linux
Updated CpuInfo::Impl in cpu_info.cc to use GetAffinityCpuCount() instead of raw std::thread::hardware_concurrency()
Updated ThreadPool::DefaultCapacity() to use affinity-aware core count as the fallback when no environment variables are set
Updated the corresponding test (TestGlobalThreadPool.Capacity) to match the new behavior across platforms
Added a Linux-only unit test (CpuInfoTest.CpuAffinity) to validate that CpuInfo reflects sched_getaffinity()
Used #ifdef __linux__ to ensure cross-platform compatibility

Are these changes tested?

CpuInfoTest.CpuAffinity checks the affinity-aware logic against sched_getaffinity() on Linux.
TestGlobalThreadPool.Capacity was updated to test the fallback behavior and env var overrides under the new logic.

Are there any user-facing changes?

No

GitHub Issue: [C++] Default thread count should count process affinity into account #45860

github-actions · 2025-04-04T19:03:28Z

Thanks for opening a pull request!

If this is not a minor PR. Could you open an issue for this pull request on GitHub? https://github.com/apache/arrow/issues/new/choose

Opening GitHub issues ahead of time contributes to the Openness of the Apache Arrow project.

Then could you also rename the pull request title in the following format?

GH-${GITHUB_ISSUE_ID}: [${COMPONENT}] ${SUMMARY}

or

MINOR: [${COMPONENT}] ${SUMMARY}

See also:

github-actions · 2025-04-04T19:04:35Z

⚠️ GitHub issue #45860 has been automatically assigned in GitHub to PR creator.

kou · 2025-04-07T08:32:30Z

cpp/src/arrow/util/affinity.h

+inline int GetAffinityCpuCount() {
+#ifdef __linux__
+  cpu_set_t mask;
+  if (sched_getaffinity(0, sizeof(mask), &mask) == 0) {
+    return CPU_COUNT(&mask);
+  }
+#endif
+  return std::thread::hardware_concurrency();
+}


How about moving this to cpp/src/arrow/util/cpu_info.cc?

kou · 2025-04-07T08:39:30Z

cpp/src/arrow/util/cpu_info.cc

    OsRetrieveCpuInfo(&hardware_flags, &vendor, &model_name);
    original_hardware_flags = hardware_flags;
-    num_cores = std::max(static_cast<int>(std::thread::hardware_concurrency()), 1);
+    num_cores = std::max(GetAffinityCpuCount(), 1);


How about not changing this and adding a new method (num_affinity_cores()?) instead?

Got it. Just to double-ceck: do we want ThreadPool::DefaultCapacity() to continue using hardware_concurrency() even on Linux, or would it make sense to prefer num_affinity_cores() there

My suggestion:

CpuInfo::num_cores() uses hardware_concurrency()

ThreadPool::DefaultCapacity() uses CpuInfo::num_affinity_cores() not CpuInfo::num_cores()

@pitrou @cyb70289 Do you have any opinion?

Looks good to me. num_cores() as hardware limit. num_affinity_cores() as actual resource available.

kou · 2025-04-07T08:41:11Z

cpp/src/arrow/util/thread_pool_test.cc

+  int expected_capacity = hw_capacity;
+#endif
+
+  ASSERT_EQ(ThreadPool::DefaultCapacity(), expected_capacity);


How about this?

Suggested change

ASSERT_EQ(ThreadPool::DefaultCapacity(), expected_capacity);

ASSERT_LT(ThreadPool::DefaultCapacity(), hw_capacity);

cyb70289 · 2025-04-08T02:35:36Z

cpp/src/arrow/util/cpu_info_test.cc

+#include <sched.h>
+#endif
+
+TEST(CpuInfoTest, CpuAffinity) {


Looks this test will never fail? But I don't have better suggestion :)

cyb70289 · 2025-04-08T02:43:13Z

cpp/src/arrow/util/thread_pool_test.cc

+#ifdef __linux__
+  ASSERT_EQ(ThreadPool::DefaultCapacity(), std::min(999, arrow::internal::GetAffinityCpuCount()));
+#else
+  ASSERT_EQ(ThreadPool::DefaultCapacity(), std::min(999, hw_capacity));
+#endif


Can we simplify it? Leave only line 1065, and remove hw_capacity?

pitrou · 2025-04-08T10:36:42Z

For the record, I don't think this should be in CpuInfo. We generally put various platform abstractions in arrow/util/io_util.h. Also, in the future we ideally want to replace CpuInfo with a third-party library, which is easier if we don't add more functionality to it. Le 7 avril 2025 09:59:15 GMT+02:00, "Raúl Cumplido" ***@***.***> a écrit :

…

@raulcd requested your review on: apache/arrow#46034 GH-45860: [C++] Respect CPU affinity in cpu_count and ThreadPool default capacity. -- Reply to this email directly or view it on GitHub: #46034 (comment) You are receiving this because your review was requested. Message ID: ***@***.***>

kou · 2025-04-09T02:00:28Z

We generally put various platform abstractions in arrow/util/io_util.h.

Oh, we have GetThreadId(), GetCurrentRSS(), ... in arrow/util/io_util.h... I feel that it's strange because it's not related to IO.

Also, in the future we ideally want to replace CpuInfo with a third-party library, which is easier if we don't add more functionality to it.

I think that we can wrap a third-party library and implement missing features by ourselves (or mix multiple third-party libraries).

pitrou · 2025-04-16T08:19:50Z

I think that we can wrap a third-party library and implement missing features by ourselves (or mix multiple third-party libraries).

That's what we did originally for CpuInfo and it means we can't easily backport improvements from upstream. So I'd rather vendor the third-party library as is and implement other functionality in other headers.

kou · 2025-04-22T05:55:00Z

I think that we can wrap a third-party library and implement missing features by ourselves (or mix multiple third-party libraries).

That's what we did originally for CpuInfo and it means we can't easily backport improvements from upstream. So I'd rather vendor the third-party library as is and implement other functionality in other headers.

You mean #785 that imports CpuInfo from parquet-cpp, right?

My idea is "wrap" not "import". For example, if we use https://github.com/google/cpu_features and https://github.com/anrieff/libcpuid , our CpuInfo will look like the following:

#include <cpuinfo_x86.h>
#include <libcpuid.h>

struct CpuInfo::Impl {
  Impl() {
    if (cpu_features::GetX86Info().features.avx) {
      hardware_flags |= CpuInfo::AVX;
    }

    struct cpu_raw_data_t raw; 
    struct cpu_id_t data;
    cpuid_get_raw_data(&raw);
    cpu_identify(&raw, &data)
    num_cores = data.num_cores;
  }
};

pitrou · 2025-04-22T07:00:09Z

My idea is "wrap" not "import". For example, if we use https://github.com/google/cpu_features and https://github.com/anrieff/libcpuid , our CpuInfo will look like the following:

Ah, I see. Yes, I agree with that.

…ult capacity (#47152) ### Rationale for this change We want the ThreadPool default capacity to follow the CPU affinity set by the user, if any. For example: ```console $ python -c "import pyarrow as pa; print(pa.cpu_count())" 24 $ taskset -c 5,6,7 python -c "import pyarrow as pa; print(pa.cpu_count())" 3 ``` ### What changes are included in this PR? - Implement and expose CPU affinity detection as a utility function in `arrow/io_util.h`; on non-Linux platform, it returns `Status::NotImplemented` - Use CPU affinity count, if available, to choose the default ThreadPool capacity (note: based on original changes by Zihan Qi in PR #46034) ### Are these changes tested? By unit tests on CI, and by hand locally. ### Are there any user-facing changes? ThreadPool capacity now follows CPU affinity settings on Linux. * GitHub Issue: #45860 Lead-authored-by: AntoinePrv <AntoinePrv@users.noreply.github.com> Co-authored-by: Zihan Qi <zihan.qi@tum.de> Signed-off-by: Antoine Pitrou <antoine@python.org>

github-actions bot added Component: C++ awaiting review Awaiting review labels Apr 4, 2025

MuteBardTison changed the title ~~ARROW-45860: Make cpu_count respect Linux CPU affinity~~ GH-45860: [C++] Make cpu_count respect Linux CPU affinity Apr 4, 2025

MuteBardTison mentioned this pull request Apr 4, 2025

[C++] Default thread count should count process affinity into account #45860

Closed

ARROW-45860: Make cpu_count respect Linux CPU affinity

d90b17e

MuteBardTison force-pushed the fix-cpu-affinity branch from 176a6fd to d90b17e Compare April 4, 2025 19:47

MuteBardTison changed the title ~~GH-45860: [C++] Make cpu_count respect Linux CPU affinity~~ GH-45860: [C++] Respect CPU affinity in cpu_count and ThreadPool default capacity Apr 4, 2025

raulcd requested review from cyb70289 and pitrou April 7, 2025 07:58

kou reviewed Apr 7, 2025

View reviewed changes

cyb70289 reviewed Apr 8, 2025

View reviewed changes

AntoinePrv mentioned this pull request Jul 21, 2025

GH-45860: [C++] Respect CPU affinity in cpu_count and ThreadPool default capacity #47152

Merged

	ASSERT_EQ(ThreadPool::DefaultCapacity(), expected_capacity);
	ASSERT_LT(ThreadPool::DefaultCapacity(), hw_capacity);

GH-45860: [C++] Respect CPU affinity in cpu_count and ThreadPool default capacity #46034

Are you sure you want to change the base?

GH-45860: [C++] Respect CPU affinity in cpu_count and ThreadPool default capacity #46034

Conversation

MuteBardTison commented Apr 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

github-actions bot commented Apr 4, 2025

Uh oh!

github-actions bot commented Apr 4, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pitrou commented Apr 8, 2025 via email

Uh oh!

kou commented Apr 9, 2025

Uh oh!

pitrou commented Apr 16, 2025

Uh oh!

kou commented Apr 22, 2025

Uh oh!

pitrou commented Apr 22, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

MuteBardTison commented Apr 4, 2025 •

edited

Loading