[coll] Improve column split tests with named threads. (#10735)

src/collective/loop.cc

@@ -14,6 +14,7 @@
 #include <thread>     // for thread
 #include <utility>    // 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

src/collective/tracker.cc

@@ -23,6 +23,7 @@
 #include <utility>    // 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<WorkerProxy>* 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) {

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 <fstream>     // for ifstream
 #include <string>      // for string
 
-#include "common.h"  // for DivRoundUp
+#include "common.h"           // for DivRoundUp
+
+#if defined(__linux__)
+#include <pthread.h>
+#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

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 <cstddef>      // for size_t
 #include <cstdint>      // for int32_t
 #include <cstdlib>      // for malloc, free
-#include <functional>   // for function
 #include <new>          // for bad_alloc
+#include <thread>       // for thread
 #include <type_traits>  // for is_signed, conditional_t, is_integral_v, invoke_result_t
 #include <vector>       // 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_

src/common/threadpool.h

@@ -9,11 +9,15 @@
 #include <memory>              // for make_shared
 #include <mutex>               // for mutex, unique_lock
 #include <queue>               // for queue
+#include <string>              // for string
 #include <thread>              // for thread
 #include <type_traits>         // for invoke_result_t
 #include <utility>             // for move
 #include <vector>              // 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 <typename InitFn>
-  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<InitFn>(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);
     }
   }
 

src/data/sparse_page_source.h

@@ -336,7 +336,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S>, public FormatStreamPol
  public:
   SparsePageSourceImpl(float missing, int nthreads, bst_feature_t n_features, bst_idx_t n_batches,
                        std::shared_ptr<Cache> 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},