diff --git a/include/xgboost/windefs.h b/include/xgboost/windefs.h index e7e743184..99bf11d09 100644 --- a/include/xgboost/windefs.h +++ b/include/xgboost/windefs.h @@ -30,4 +30,4 @@ #endif // xgboost_IS_MINGW -#endif // defined(xgboost_IS_WIN) +#endif // !defined(xgboost_IS_WIN) diff --git a/src/collective/loop.cc b/src/collective/loop.cc index cd7740fd2..6949ebc53 100644 --- a/src/collective/loop.cc +++ b/src/collective/loop.cc @@ -14,6 +14,7 @@ #include // for thread #include // for move +#include "../common/threading_utils.h" // for NameThread #include "xgboost/collective/poll_utils.h" // for PollHelper #include "xgboost/collective/result.h" // for Fail, Success #include "xgboost/collective/socket.h" // for FailWithCode @@ -271,5 +272,6 @@ Loop::Loop(std::chrono::seconds timeout) : timeout_{timeout} { worker_ = std::thread{[this] { this->Process(); }}; + common::NameThread(&worker_, "lw"); } } // namespace xgboost::collective diff --git a/src/collective/tracker.cc b/src/collective/tracker.cc index bbc7a7c5a..b1081fe8e 100644 --- a/src/collective/tracker.cc +++ b/src/collective/tracker.cc @@ -23,6 +23,7 @@ #include // for move, forward #include "../common/json_utils.h" +#include "../common/threading_utils.h" // for NameThread #include "comm.h" #include "protocol.h" // for kMagic, PeerInfo #include "tracker.h" @@ -143,6 +144,8 @@ Result RabitTracker::Bootstrap(std::vector* p_workers) { Json::Dump(jnext, &str); worker.Send(StringView{str}); }); + std::string name = "tkbs_t-" + std::to_string(r); + common::NameThread(&bootstrap_threads.back(), name.c_str()); } for (auto& t : bootstrap_threads) { diff --git a/src/common/threading_utils.cc b/src/common/threading_utils.cc index 1f4d5be2f..625c98d1c 100644 --- a/src/common/threading_utils.cc +++ b/src/common/threading_utils.cc @@ -1,5 +1,5 @@ /** - * Copyright 2022-2023 by XGBoost Contributors + * Copyright 2022-2024, XGBoost Contributors */ #include "threading_utils.h" @@ -9,7 +9,11 @@ #include // for ifstream #include // for string -#include "common.h" // for DivRoundUp +#include "common.h" // for DivRoundUp + +#if defined(__linux__) +#include +#endif namespace xgboost::common { /** @@ -113,4 +117,26 @@ std::int32_t OmpGetNumThreads(std::int32_t n_threads) { n_threads = std::max(n_threads, 1); return n_threads; } + +void NameThread(std::thread* t, StringView name) { +#if defined(__linux__) + auto handle = t->native_handle(); + char old[16]; + auto ret = pthread_getname_np(handle, old, 16); + if (ret != 0) { + LOG(WARNING) << "Failed to get the name from thread"; + } + auto new_name = std::string{old} + ">" + name.c_str(); // NOLINT + if (new_name.size() > 15) { + new_name = new_name.substr(new_name.size() - 15); + } + ret = pthread_setname_np(handle, new_name.c_str()); + if (ret != 0) { + LOG(WARNING) << "Failed to name thread:" << ret << " :" << new_name; + } +#else + (void)name; + (void)t; +#endif +} } // namespace xgboost::common diff --git a/src/common/threading_utils.h b/src/common/threading_utils.h index ac7119035..e21400705 100644 --- a/src/common/threading_utils.h +++ b/src/common/threading_utils.h @@ -1,5 +1,5 @@ /** - * Copyright 2019-2023 by XGBoost Contributors + * Copyright 2019-2024, XGBoost Contributors */ #ifndef XGBOOST_COMMON_THREADING_UTILS_H_ #define XGBOOST_COMMON_THREADING_UTILS_H_ @@ -11,12 +11,13 @@ #include // for size_t #include // for int32_t #include // for malloc, free -#include // for function #include // for bad_alloc +#include // for thread #include // for is_signed, conditional_t, is_integral_v, invoke_result_t #include // for vector #include "xgboost/logging.h" +#include "xgboost/string_view.h" // for StringView #if !defined(_OPENMP) extern "C" { @@ -308,6 +309,11 @@ class MemStackAllocator { * \brief Constant that can be used for initializing static thread local memory. */ std::int32_t constexpr DefaultMaxThreads() { return 128; } + +/** + * @brief Give the thread a name. Supports only pthread on linux. + */ +void NameThread(std::thread* t, StringView name); } // namespace xgboost::common #endif // XGBOOST_COMMON_THREADING_UTILS_H_ diff --git a/src/common/threadpool.h b/src/common/threadpool.h index 21e27aa76..14469df91 100644 --- a/src/common/threadpool.h +++ b/src/common/threadpool.h @@ -9,11 +9,15 @@ #include // for make_shared #include // for mutex, unique_lock #include // for queue +#include // for string #include // for thread #include // for invoke_result_t #include // for move #include // for vector +#include "threading_utils.h" // for NameThread +#include "xgboost/string_view.h" // for StringView + namespace xgboost::common { /** * @brief Simple implementation of a thread pool. @@ -27,11 +31,12 @@ class ThreadPool { public: /** + * @param name Name prefix for threads. * @param n_threads The number of threads this pool should hold. * @param init_fn Function called once during thread creation. */ template - explicit ThreadPool(std::int32_t n_threads, InitFn&& init_fn) { + explicit ThreadPool(StringView name, std::int32_t n_threads, InitFn&& init_fn) { for (std::int32_t i = 0; i < n_threads; ++i) { pool_.emplace_back([&, init_fn = std::forward(init_fn)] { init_fn(); @@ -55,6 +60,8 @@ class ThreadPool { fn(); } }); + std::string name_i = name.c_str() + std::string{"-"} + std::to_string(i); // NOLINT + NameThread(&pool_.back(), name_i); } } diff --git a/src/data/sparse_page_source.h b/src/data/sparse_page_source.h index ca04e969f..5ec93a311 100644 --- a/src/data/sparse_page_source.h +++ b/src/data/sparse_page_source.h @@ -336,7 +336,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl, public FormatStreamPol public: SparsePageSourceImpl(float missing, int nthreads, bst_feature_t n_features, bst_idx_t n_batches, std::shared_ptr cache) - : workers_{std::max(2, std::min(nthreads, 16)), InitNewThread{}}, + : workers_{StringView{"ext-mem"}, std::max(2, std::min(nthreads, 16)), InitNewThread{}}, missing_{missing}, nthreads_{nthreads}, n_features_{n_features}, diff --git a/tests/cpp/collective/test_worker.h b/tests/cpp/collective/test_worker.h index 631a225b3..2740190f0 100644 --- a/tests/cpp/collective/test_worker.h +++ b/tests/cpp/collective/test_worker.h @@ -12,10 +12,11 @@ #include // for move #include // for vector -#include "../../../src/collective/comm.h" +#include "../../../src/collective/comm.h" // for RabitComm #include "../../../src/collective/communicator-inl.h" // for Init, Finalize #include "../../../src/collective/tracker.h" // for GetHostAddress #include "../../../src/common/cuda_rt_utils.h" // for AllVisibleGPUs +#include "../../../src/common/threading_utils.h" // for NameThread #include "../helpers.h" // for FileExists #if defined(XGBOOST_USE_FEDERATED) @@ -176,6 +177,9 @@ void TestDistributedGlobal(std::int32_t n_workers, WorkerFn worker_fn, bool need CHECK(status == std::future_status::ready) << "Test timeout"; fut.get(); }); + + std::string name = "tw-" + std::to_string(i); + common::NameThread(&workers.back(), name.c_str()); } for (auto& t : workers) { @@ -199,7 +203,7 @@ class BaseMGPUTest : public ::testing::Test { * available. */ template - auto DoTest(Fn&& fn, bool is_federated, bool emulate_if_single = false) const { + auto DoTest(Fn&& fn, bool is_federated, [[maybe_unused]] bool emulate_if_single = false) const { auto n_gpus = common::AllVisibleGPUs(); if (is_federated) { #if defined(XGBOOST_USE_FEDERATED) diff --git a/tests/cpp/common/test_threadpool.cc b/tests/cpp/common/test_threadpool.cc index ca8a73b55..5d863689b 100644 --- a/tests/cpp/common/test_threadpool.cc +++ b/tests/cpp/common/test_threadpool.cc @@ -21,7 +21,7 @@ TEST(ThreadPool, Basic) { // 4 is an invalid value, it's only possible to set it by bypassing the parameter // validation. ASSERT_NE(orig, GlobalConfigThreadLocalStore::Get()->verbosity); - ThreadPool pool{n_threads, [config = *GlobalConfigThreadLocalStore::Get()] { + ThreadPool pool{StringView{"test"}, n_threads, [config = *GlobalConfigThreadLocalStore::Get()] { *GlobalConfigThreadLocalStore::Get() = config; }}; GlobalConfigThreadLocalStore::Get()->verbosity = orig; // restore diff --git a/tests/cpp/test_learner.cc b/tests/cpp/test_learner.cc index d53a568d4..58e52e63d 100644 --- a/tests/cpp/test_learner.cc +++ b/tests/cpp/test_learner.cc @@ -745,8 +745,7 @@ void VerifyColumnSplitWithArgs(std::string const& tree_method, bool use_gpu, Arg std::shared_ptr sliced{p_fmat->SliceCol(world_size, rank)}; std::string device = "cpu"; if (use_gpu) { - auto gpu_id = common::AllVisibleGPUs() == 1 ? 0 : rank; - device = "cuda:" + std::to_string(gpu_id); + device = MakeCUDACtx(DistGpuIdx()).DeviceName(); } auto model = GetModelWithArgs(sliced, tree_method, device, args); ASSERT_EQ(model, expected_model); @@ -807,44 +806,32 @@ class ColumnSplitTrainingTest } }; -auto MakeParamsForTest() { - std::vector> configs; - for (auto tm : {"hist", "approx"}) { -#if defined(XGBOOST_USE_CUDA) - std::array use_gpu{true, false}; -#else - std::array use_gpu{false}; -#endif - for (auto i : use_gpu) { +auto WithFed() { #if defined(XGBOOST_USE_FEDERATED) - std::array fed{true, false}; + return ::testing::Bool(); #else - std::array fed{false}; + return ::testing::Values(false); #endif - for (auto j : fed) { - configs.emplace_back(tm, i, j); - } - } - } - return configs; } } // anonymous namespace TEST_P(ColumnSplitTrainingTest, ColumnSampler) { - auto param = GetParam(); - std::apply(TestColumnSplitColumnSampler, param); + std::apply(TestColumnSplitColumnSampler, GetParam()); } TEST_P(ColumnSplitTrainingTest, InteractionConstraints) { - auto param = GetParam(); - std::apply(TestColumnSplitInteractionConstraints, param); + std::apply(TestColumnSplitInteractionConstraints, GetParam()); } TEST_P(ColumnSplitTrainingTest, MonotoneConstraints) { - auto param = GetParam(); - std::apply(TestColumnSplitMonotoneConstraints, param); + std::apply(TestColumnSplitMonotoneConstraints, GetParam()); } -INSTANTIATE_TEST_SUITE_P(ColumnSplit, ColumnSplitTrainingTest, - ::testing::ValuesIn(MakeParamsForTest())); +INSTANTIATE_TEST_SUITE_P(Cpu, ColumnSplitTrainingTest, + ::testing::Combine(::testing::Values("hist", "approx"), + ::testing::Values(false), WithFed())); + +INSTANTIATE_TEST_SUITE_P(MGPU, ColumnSplitTrainingTest, + ::testing::Combine(::testing::Values("hist", "approx"), + ::testing::Values(true), WithFed())); } // namespace xgboost