Support CUDA ordinal for external memory. (dmlc#11219)

Author: trivialfis

include/xgboost/global_config.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020-2024, XGBoost Contributors
+ * Copyright 2020-2025, XGBoost Contributors
  * \file global_config.h
  * \brief Global configuration for XGBoost
  * \author Hyunsu Cho
@@ -31,9 +31,11 @@ struct GlobalConfiguration : public XGBoostParameter<GlobalConfiguration> {
 using GlobalConfigThreadLocalStore = dmlc::ThreadLocalStore<GlobalConfiguration>;
 
 struct InitNewThread {
-  GlobalConfiguration config = *GlobalConfigThreadLocalStore::Get();
+  GlobalConfiguration config;
+  std::int32_t device{-1};
 
   void operator()() const;
+  InitNewThread();
 };
 }  // namespace xgboost
 

src/common/cuda_rt_utils.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2015-2024, XGBoost Contributors
+ * Copyright 2015-2025, XGBoost Contributors
  */
 #include "cuda_rt_utils.h"
 
@@ -28,9 +28,14 @@ std::int32_t AllVisibleGPUs() {
   return n_visgpus;
 }
 
-std::int32_t CurrentDevice() {
-  std::int32_t device = 0;
-  dh::safe_cuda(cudaGetDevice(&device));
+std::int32_t CurrentDevice(bool raise) {
+  std::int32_t device = -1;
+  if (raise) {
+    dh::safe_cuda(cudaGetDevice(&device));
+  } else if (cudaGetDevice(&device) != cudaSuccess) {
+    // Return -1 as an error.
+    return -1;
+  }
   return device;
 }
 
@@ -100,8 +105,10 @@ void DrVersion(std::int32_t* major, std::int32_t* minor) {
 #else
 std::int32_t AllVisibleGPUs() { return 0; }
 
-std::int32_t CurrentDevice() {
-  common::AssertGPUSupport();
+std::int32_t CurrentDevice(bool raise) {
+  if (raise) {
+    common::AssertGPUSupport();
+  }
   return -1;
 }
 

src/common/cuda_rt_utils.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2024, XGBoost contributors
+ * Copyright 2024-2025, XGBoost contributors
  */
 #pragma once
 #include <cstddef>  // for size_t
@@ -12,7 +12,10 @@
 namespace xgboost::curt {
 std::int32_t AllVisibleGPUs();
 
-std::int32_t CurrentDevice();
+/**
+ * @param raise Raise error if XGBoost is not compiled with CUDA, or GPU is not available.
+ */
+std::int32_t CurrentDevice(bool raise = true);
 
 // Whether the device supports coherently accessing pageable memory without calling
 // `cudaHostRegister` on it

src/data/ellpack_page_source.cu

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019-2024, XGBoost contributors
+ * Copyright 2019-2025, XGBoost contributors
  */
 #include <algorithm>  // for count_if
 #include <cstddef>    // for size_t
@@ -8,8 +8,7 @@
 #include <numeric>    // for accumulate
 #include <utility>    // for move
 
-#include "../common/common.h"               // for safe_cuda
-#include "../common/common.h"               // for HumanMemUnit
+#include "../common/common.h"               // for HumanMemUnit, safe_cuda
 #include "../common/cuda_rt_utils.h"        // for SetDevice
 #include "../common/device_helpers.cuh"     // for CUDAStreamView, DefaultStream
 #include "../common/ref_resource_view.cuh"  // for MakeFixedVecWithCudaMalloc

src/global_config.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2020-2024, XGBoost Contributors
+ * Copyright 2020-2025, XGBoost Contributors
  * \file global_config.cc
  * \brief Global configuration for XGBoost
  * \author Hyunsu Cho
@@ -9,13 +9,21 @@
 
 #include <dmlc/thread_local.h>
 
+#include "common/cuda_rt_utils.h"  // for SetDevice
+
 namespace xgboost {
 DMLC_REGISTER_PARAMETER(GlobalConfiguration);
 
+InitNewThread::InitNewThread()
+    : config{*GlobalConfigThreadLocalStore::Get()}, device{curt::CurrentDevice(false)} {}
+
 void InitNewThread::operator()() const {
   *GlobalConfigThreadLocalStore::Get() = config;
   if (config.nthread > 0) {
     omp_set_num_threads(config.nthread);
   }
+  if (device >= 0) {
+    curt::SetDevice(this->device);
+  }
 }
 }  // namespace xgboost

tests/test_distributed/test_gpu_with_dask/test_gpu_external_memory.py

@@ -1,10 +1,17 @@
-"""Copyright 2024, XGBoost contributors"""
+"""Copyright 2024-2025, XGBoost contributors"""
+
+from functools import partial, update_wrapper
+from typing import Any
 
 import pytest
 from dask_cuda import LocalCUDACluster
 from distributed import Client
 
+import xgboost as xgb
+from xgboost import collective as coll
+from xgboost import testing as tm
 from xgboost.testing.dask import check_external_memory, get_rabit_args
+from xgboost.tracker import RabitTracker
 
 
 @pytest.mark.parametrize("is_qdm", [True, False])
@@ -22,3 +29,48 @@ def test_external_memory(is_qdm: bool) -> None:
                 is_qdm=is_qdm,
             )
             client.gather(futs)
+
+
+@pytest.mark.skipif(**tm.no_loky())
+def test_extmem_qdm_distributed() -> None:
+    from loky import get_reusable_executor
+
+    n_samples_per_batch = 2048
+    n_features = 128
+    n_batches = 8
+
+    def do_train(ordinal: int) -> None:
+        it = tm.IteratorForTest(
+            *tm.make_batches(n_samples_per_batch, n_features, n_batches, use_cupy=True),
+            cache="cache",
+            on_host=True,
+        )
+
+        Xy = xgb.ExtMemQuantileDMatrix(it)
+        results: dict[str, Any] = {}
+        booster = xgb.train(
+            {"device": f"cuda:{ordinal}"},
+            num_boost_round=2,
+            dtrain=Xy,
+            evals=[(Xy, "Train")],
+            evals_result=results,
+        )
+        assert tm.non_increasing(results["Train"]["rmse"])
+
+    tracker = RabitTracker(host_ip="127.0.0.1", n_workers=2)
+    tracker.start()
+    args = tracker.worker_args()
+
+    def local_test(worker_id: int, rabit_args: dict) -> None:
+        import cupy as cp
+
+        cp.cuda.runtime.setDevice(worker_id)
+
+        with coll.CommunicatorContext(**rabit_args, DMLC_TASK_ID=str(worker_id)):
+            assert coll.get_rank() == worker_id
+            do_train(coll.get_rank())
+
+    n_workers = 2
+    fn = update_wrapper(partial(local_test, rabit_args=args), local_test)
+    with get_reusable_executor(max_workers=n_workers) as pool:
+        results = pool.map(fn, range(n_workers))