Reduce large allocation tests to 70% on WSL and probe max allocatable memory

Author: bdice

cpp/tests/mr/arena_mr_tests.cpp

@@ -4,6 +4,7 @@
  */
 
 #include "../byte_literals.hpp"
+#include "test_utils.hpp"
 
 #include <rmm/aligned.hpp>
 #include <rmm/cuda_device.hpp>
@@ -483,11 +484,13 @@ TEST_F(ArenaTest, SizeSmallerThanSuperblockSize)  // NOLINT
   EXPECT_THROW(construct_small(), rmm::logic_error);
 }
 
-TEST_F(ArenaTest, AllocateNinetyPercent)  // NOLINT
+TEST_F(ArenaTest, AllocateMostOfFreeMemory)  // NOLINT
 {
-  EXPECT_NO_THROW([]() {  // NOLINT(cppcoreguidelines-avoid-goto)
-    auto const ninety_percent = rmm::percent_of_free_device_memory(90);
-    arena_mr mr(rmm::mr::get_current_device_resource_ref(), ninety_percent);
+  if (is_wsl()) { log_max_allocatable_memory(); }
+  auto const percent = is_wsl() ? 70 : 90;
+  EXPECT_NO_THROW([percent]() {  // NOLINT(cppcoreguidelines-avoid-goto)
+    auto const pool_size = rmm::percent_of_free_device_memory(percent);
+    arena_mr mr(rmm::mr::get_current_device_resource_ref(), pool_size);
   }());
 }
 

cpp/tests/mr/pool_mr_tests.cpp

@@ -3,6 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include "test_utils.hpp"
+
 #include <rmm/cuda_device.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/device_buffer.hpp>
@@ -45,15 +47,15 @@ TEST(PoolTest, ThrowMaxLessThanInitial)
   EXPECT_THROW(max_less_than_initial(), rmm::logic_error);
 }
 
-TEST(PoolTest, AllocateNinetyPercent)
+TEST(PoolTest, AllocateMostOfFreeMemory)
 {
-  auto allocate_ninety = []() {
-    auto const [free, total] = rmm::available_device_memory();
-    (void)total;
-    auto const ninety_percent_pool = rmm::percent_of_free_device_memory(90);
-    pool_mr mr{rmm::mr::get_current_device_resource_ref(), ninety_percent_pool};
+  if (is_wsl()) { log_max_allocatable_memory(); }
+  auto const percent = is_wsl() ? 70 : 90;
+  auto allocate      = [percent]() {
+    auto const pool_size = rmm::percent_of_free_device_memory(percent);
+    pool_mr mr{rmm::mr::get_current_device_resource_ref(), pool_size};
   };
-  EXPECT_NO_THROW(allocate_ninety());
+  EXPECT_NO_THROW(allocate());
 }
 
 TEST(PoolTest, TwoLargeBuffers)

cpp/tests/mr/test_utils.hpp

@@ -1,17 +1,67 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
 #include <rmm/aligned.hpp>
+#include <rmm/cuda_device.hpp>
+#include <rmm/detail/format.hpp>
 #include <rmm/mr/system_memory_resource.hpp>
 
 #include <cuda_runtime_api.h>
 
+#include <cstddef>
+#include <fstream>
+#include <iostream>
+#include <string>
+
 namespace rmm::test {
 
+/**
+ * @brief Returns true if running under Windows Subsystem for Linux (WSL).
+ */
+inline bool is_wsl()
+{
+  std::ifstream proc_version("/proc/version");
+  if (proc_version.is_open()) {
+    std::string line;
+    std::getline(proc_version, line);
+    return line.find("microsoft") != std::string::npos ||
+           line.find("Microsoft") != std::string::npos;
+  }
+  return false;
+}
+
+/**
+ * @brief Probes and logs the maximum allocatable GPU memory percentage.
+ *
+ * Tries 90%, 80%, 70%, 60% of free memory in sequence, stopping at the first
+ * successful cudaMalloc. Logs the result along with cudaMemGetInfo values.
+ */
+inline void log_max_allocatable_memory()
+{
+  auto const [free, total] = rmm::available_device_memory();
+  std::cout << "cudaMemGetInfo: free=" << rmm::detail::format_bytes(free)
+            << " total=" << rmm::detail::format_bytes(total) << std::endl;
+
+  for (int pct : {90, 80, 70, 60}) {
+    auto const size = static_cast<std::size_t>(static_cast<double>(free) * pct / 100.0);
+    void* ptr       = nullptr;
+    if (cudaMalloc(&ptr, size) == cudaSuccess) {
+      cudaFree(ptr);
+      std::cout << "Max allocatable: >= " << pct << "% of free (" << rmm::detail::format_bytes(size)
+                << ")" << std::endl;
+      return;
+    }
+    cudaGetLastError();
+    std::cout << "cudaMalloc failed at " << pct << "% of free (" << rmm::detail::format_bytes(size)
+              << ")" << std::endl;
+  }
+  std::cout << "cudaMalloc failed at all tested percentages" << std::endl;
+}
+
 /**
  * @brief Returns if a pointer points to a device memory or managed memory
  * allocation.