[backport] Allow fallback for CUDA virtual memory. (dmlc#11391, dmlc#11398) (dmlc#11434)

- Check CUDA version based on both driver API and nvidia-smi.

Author: trivialfis

src/common/common.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2015-2024, XGBoost Contributors
+ * Copyright 2015-2025, XGBoost Contributors
  * \file common.h
  * \brief Common utilities
  */
@@ -62,6 +62,19 @@ inline std::vector<std::string> Split(const std::string& s, char delim) {
   return ret;
 }
 
+// Trims leading whitespace from a string
+[[nodiscard]] inline std::string TrimFirst(const std::string &str) {
+  if (str.empty()) {
+    return str;
+  }
+
+  std::size_t first = str.find_first_not_of(" \t\n\r");
+  if (first == std::string::npos) {
+    return "";
+  }
+  return str.substr(first);
+}
+
 /**
  * @brief Add escapes for a UTF-8 string.
  */

src/common/cuda_dr_utils.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2024, XGBoost contributors
+ * Copyright 2024-2025, XGBoost contributors
  */
 #if defined(XGBOOST_USE_CUDA)
 #include "cuda_dr_utils.h"
@@ -10,11 +10,12 @@
 #include <memory>     // for make_unique
 #include <mutex>      // for call_once
 #include <sstream>    // for stringstream
-#include <string>     // for string
+#include <string>     // for string, stoi
 
-#include "common.h"               // for safe_cuda
+#include "common.h"               // for safe_cuda, TrimFirst, Split
 #include "cuda_rt_utils.h"        // for CurrentDevice
-#include "xgboost/string_view.h"  // for StringVie
+#include "io.h"                   // for CmdOutput
+#include "xgboost/string_view.h"  // for StringView
 
 namespace xgboost::cudr {
 CuDriverApi::CuDriverApi() {
@@ -104,5 +105,52 @@ void MakeCuMemLocation(CUmemLocationType type, CUmemLocation *loc) {
   MakeCuMemLocation(type, &prop.location);
   return prop;
 }
+
+[[nodiscard]] bool GetVersionFromSmi(std::int32_t *p_major, std::int32_t *p_minor) {
+  using ::xgboost::common::Split;
+  using ::xgboost::common::TrimFirst;
+  // `nvidia-smi --version` is not available for older versions, as a result, we can't query the
+  // cuda driver version unless we want to parse the table output.
+
+  // Example output on a 2-GPU system:
+  //
+  // $ nvidia-smi --query-gpu=driver_version --format=csv
+  //
+  // driver_version
+  // 570.124.06
+  // 570.124.06
+  //
+  auto cmd = "nvidia-smi --query-gpu=driver_version --format=csv";
+  auto smi_out_str = common::CmdOutput(StringView{cmd});
+
+  auto Invalid = [=] {
+    *p_major = *p_minor = -1;
+    return false;
+  };
+  if (smi_out_str.empty()) {
+    return Invalid();
+  }
+
+  auto smi_split = Split(smi_out_str, '\n');
+  if (smi_split.size() < 2) {
+    return Invalid();
+  }
+
+  // Use the first GPU
+  auto smi_ver = Split(TrimFirst(smi_split[1]), '.');
+  // 570.124.06
+  if (smi_ver.size() != 3) {
+    return Invalid();
+  }
+  try {
+    *p_major = std::stoi(smi_ver[0]);
+    *p_minor = std::stoi(smi_ver[1]);
+    LOG(INFO) << "Driver version: `" << *p_major << "." << *p_minor << "`";
+    return true;
+  } catch (std::exception const &) {
+  }
+
+  return Invalid();
+}
 }  // namespace xgboost::cudr
 #endif

src/common/cuda_dr_utils.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2024, XGBoost contributors
+ * Copyright 2024-2025, XGBoost contributors
  *
  * @brief Utility for CUDA driver API.
  *
@@ -102,4 +102,11 @@ void MakeCuMemLocation(CUmemLocationType type, CUmemLocation* loc);
  * @brief Construct a `CUmemAllocationProp`.
  */
 [[nodiscard]] CUmemAllocationProp MakeAllocProp(CUmemLocationType type);
+
+/**
+ * @brief Get system driver version from the `nvidia-smi` command.
+ *
+ * @return Whether the system call is successful.
+ */
+[[nodiscard]] bool GetVersionFromSmi(std::int32_t *p_major, std::int32_t *p_minor);
 }  // namespace xgboost::cudr

src/common/device_helpers.cu

@@ -1,19 +1,51 @@
 /**
- * Copyright 2024, XGBoost contributors
+ * Copyright 2024-2025, XGBoost contributors
  */
-#include "cuda_rt_utils.h"  // for RtVersion
+#include <mutex>  // for once_flag, call_once
+
+#include "../common/cuda_dr_utils.h"  // for GetVersionFromSmi
+#include "cuda_rt_utils.h"            // for RtVersion
 #include "device_helpers.cuh"
+#include "device_vector.cuh"  // for GrowOnlyVirtualMemVec
 #include "xgboost/windefs.h"  // for xgboost_IS_WIN
 
 namespace dh {
+namespace {
+[[nodiscard]] bool IsSupportedDrVer(std::int32_t major, std::int32_t minor) {
+  return major > 12 || (major == 12 && minor >= 5);
+}
+
+// Check whether cuda virtual memory can be used.
+// Host NUMA allocation requires driver that supports CTK >= 12.5 to be stable
+[[nodiscard]] bool CheckVmAlloc() {
+  static bool vm_flag = true;
+  static std::once_flag once;
+
+  std::call_once(once, [] {
+    std::int32_t major{0}, minor{0};
+    xgboost::curt::DrVersion(&major, &minor);
+    if (IsSupportedDrVer(major, minor)) {
+      // The result from the driver api is not reliable. The system driver might not match
+      // the CUDA driver in some obscure cases.
+      //
+      // https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
+      // Ver                 Linux       Win
+      // CUDA 12.5 Update 1  >=555.42.06 >=555.85
+      // CUDA 12.5 GA        >=555.42.02 >=555.85
+      vm_flag = xgboost::cudr::GetVersionFromSmi(&major, &minor) && major >= 555;
+    } else {
+      vm_flag = false;
+    }
+  });
+  return vm_flag;
+}
+}  // namespace
+
 PinnedMemory::PinnedMemory() {
 #if defined(xgboost_IS_WIN)
   this->impl_.emplace<detail::GrowOnlyPinnedMemoryImpl>();
 #else
-  std::int32_t major{0}, minor{0};
-  xgboost::curt::DrVersion(&major, &minor);
-  // Host NUMA allocation requires driver that supports CTK >= 12.5 to be stable.
-  if (major >= 12 && minor >= 5) {
+  if (CheckVmAlloc()) {
     this->impl_.emplace<detail::GrowOnlyVirtualMemVec>(CU_MEM_LOCATION_TYPE_HOST_NUMA);
   } else {
     this->impl_.emplace<detail::GrowOnlyPinnedMemoryImpl>();

src/common/io.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019-2024, by XGBoost Contributors
+ * Copyright 2019-2025, by XGBoost Contributors
  */
 #if defined(__unix__) || defined(__APPLE__)
 
@@ -234,7 +234,7 @@ void detail::CloseMmap(MMAPFile* handle) {
   }
 #if defined(xgboost_IS_WIN)
   if (handle->base_ptr) {
-    CHECK(UnmapViewOfFile(handle->base_ptr)) "Faled to call munmap: " << SystemErrorMsg();
+    CHECK(UnmapViewOfFile(handle->base_ptr)) << "Failed to call munmap: " << SystemErrorMsg();
   }
   if (handle->fd != INVALID_HANDLE_VALUE) {
     CHECK(CloseHandle(handle->fd)) << "Failed to close handle: " << SystemErrorMsg();
@@ -245,11 +245,11 @@ void detail::CloseMmap(MMAPFile* handle) {
 #else
   if (handle->base_ptr) {
     CHECK_NE(munmap(handle->base_ptr, handle->base_size), -1)
-        << "Faled to call munmap: `" << handle->path << "`. " << SystemErrorMsg();
+        << "Failed to call munmap: `" << handle->path << "`. " << SystemErrorMsg();
   }
   if (handle->fd != 0) {
     CHECK_NE(close(handle->fd), -1)
-        << "Faled to close: `" << handle->path << "`. " << SystemErrorMsg();
+        << "Failed to close: `" << handle->path << "`. " << SystemErrorMsg();
   }
 #endif
   delete handle;
@@ -302,4 +302,23 @@ AlignedMemWriteStream::~AlignedMemWriteStream() = default;
 [[nodiscard]] std::size_t AlignedMemWriteStream::Tell() const noexcept(true) {
   return this->pimpl_->Tell();
 }
+
+[[nodiscard]] std::string CmdOutput(StringView cmd) {
+#if defined(xgboost_IS_WIN) || defined(__i386__)
+  (void)cmd;
+  LOG(FATAL) << "Not implemented";
+  return "";
+#else
+  // popen is a convenient method, but it always returns a success even if the command
+  // fails.
+  std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd.c_str(), "r"), pclose);
+  CHECK(pipe);
+  std::array<char, 128> buffer;
+  std::string result;
+  while (std::fgets(buffer.data(), static_cast<std::int32_t>(buffer.size()), pipe.get())) {
+    result += buffer.data();
+  }
+  return result;
+#endif
+}
 }  // namespace xgboost::common

src/common/io.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2014-2024, XGBoost Contributors
+ * Copyright 2014-2025, XGBoost Contributors
  * \file io.h
  * \brief general stream interface for serialization, I/O
  * \author Tianqi Chen
@@ -607,5 +607,8 @@ class AlignedMemWriteStream : public AlignedFileWriteStream {
 
   [[nodiscard]] std::size_t Tell() const noexcept(true);
 };
+
+// Run a system command, get its stdout.
+[[nodiscard]] std::string CmdOutput(StringView cmd);
 }  // namespace xgboost::common
 #endif  // XGBOOST_COMMON_IO_H_

tests/cpp/common/test_cuda_dr_utils.cc

@@ -0,0 +1,23 @@
+/**
+ * Copyright 2025, XGBoost Contributors
+ */
+#include <gtest/gtest.h>
+
+#if defined(XGBOOST_USE_CUDA) && defined(__linux__)
+#include "../../../src/common/cuda_dr_utils.h"
+
+namespace xgboost::cudr {
+TEST(DrUtils, GetVersionFromSmi) {
+  std::int32_t major = 0, minor = 0;
+  bool result = GetVersionFromSmi(&major, &minor);
+
+  if (result) {
+    EXPECT_GE(major, 0);
+    EXPECT_GE(minor, 0);
+  } else {
+    EXPECT_EQ(major, -1);
+    EXPECT_EQ(minor, -1);
+  }
+}
+}  // namespace xgboost::cudr
+#endif  // defined(XGBOOST_USE_CUDA)