Enable reporting peak memory usage for gtests (#18599)

Enables libcudf gtests to report peak memory usage after the tests complete.
The memory peak uses the `rmm::mr::statistics_resource_adaptor` and is triggered with the environment variable `GTEST_CUDF_MEMORY_PEAK`
Working on this uncovered that most of the `STREAM_` based tests were not using the `CUDF_TEST_PROGRAM_MAIN()` and so did not support custom parameters. Also, at least one test failed the stream check after this was corrected.
The PR includes a shell script to run each test and report the peak memory for each to stdout in CSV format.

Authors:
  - David Wendt (https://github.com/davidwendt)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Shruti Shivakumar (https://github.com/shrshi)

URL: #18599

Author: davidwendt

cpp/include/cudf_test/testing_main.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,6 +33,9 @@
 #include <rmm/mr/device/managed_memory_resource.hpp>
 #include <rmm/mr/device/owning_wrapper.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
+#include <rmm/mr/device/statistics_resource_adaptor.hpp>
+
+#include <iostream>
 
 namespace CUDF_EXPORT cudf {
 namespace test {
@@ -226,10 +229,19 @@ inline void init_cudf_test(int argc, char** argv, cudf::test::config const& conf
  * This `main` function is a wrapper around the google test generated `main`,
  * maintaining the original functionality.
  */
-#define CUDF_TEST_PROGRAM_MAIN()            \
-  int main(int argc, char** argv)           \
-  {                                         \
-    ::testing::InitGoogleTest(&argc, argv); \
-    init_cudf_test(argc, argv);             \
-    return RUN_ALL_TESTS();                 \
+#define CUDF_TEST_PROGRAM_MAIN()                                                                 \
+  int main(int argc, char** argv)                                                                \
+  {                                                                                              \
+    ::testing::InitGoogleTest(&argc, argv);                                                      \
+    init_cudf_test(argc, argv);                                                                  \
+    if (std::getenv("GTEST_CUDF_MEMORY_PEAK")) {                                                 \
+      auto mr = rmm::mr::statistics_resource_adaptor<rmm::mr::device_memory_resource>(           \
+        cudf::get_current_device_resource());                                                    \
+      cudf::set_current_device_resource(&mr);                                                    \
+      auto rc = RUN_ALL_TESTS();                                                                 \
+      std::cout << "Peak memory usage " << mr.get_bytes_counter().peak << " bytes" << std::endl; \
+      return rc;                                                                                 \
+    } else {                                                                                     \
+      return RUN_ALL_TESTS();                                                                    \
+    }                                                                                            \
   }

cpp/scripts/gtest_memory_usage.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+# Copyright (c) 2025, NVIDIA CORPORATION.
+
+export GTEST_CUDF_RMM_MODE=cuda
+export GTEST_CUDF_MEMORY_PEAK=1
+export GTEST_BRIEF=1
+for gt in gtests/*_TEST ; do
+  test_name=$(basename "${gt}")
+  echo -n "$test_name"
+  # dependent on the string output from testing_main.hpp
+  bytes=$(${gt} 2>/dev/null | grep Peak | cut -d' ' -f4)
+  echo ",${bytes}"
+done
+unset GTEST_BRIEF
+unset GTEST_CUDF_RMM_MODE
+unset GTEST_CUDF_MEMORY_PEAK

cpp/tests/CMakeLists.txt

@@ -700,22 +700,24 @@ ConfigureTest(STREAM_BINARYOP_TEST streams/binaryop_test.cpp STREAM_MODE testing
 ConfigureTest(STREAM_COLUMN_VIEW_TEST streams/column_view_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_CONCATENATE_TEST streams/concatenate_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_COPYING_TEST streams/copying_test.cpp STREAM_MODE testing)
-ConfigureTest(STREAM_CSVIO_TEST streams/io/csv_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_DATETIME_TEST streams/datetime_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_DICTIONARY_TEST streams/dictionary_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_FILLING_TEST streams/filling_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_GROUPBY_TEST streams/groupby_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_HASHING_TEST streams/hash_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_INTEROP streams/interop_test.cpp STREAM_MODE testing)
+ConfigureTest(STREAM_IO_CSV_TEST streams/io/csv_test.cpp STREAM_MODE testing)
+ConfigureTest(STREAM_IO_JSON_TEST streams/io/json_test.cpp STREAM_MODE testing)
+ConfigureTest(
+  STREAM_IO_MULTIBYTE_SPLIT_TEST streams/io/multibyte_split_test.cpp STREAM_MODE testing
+)
+ConfigureTest(STREAM_IO_ORC_TEST streams/io/orc_test.cpp STREAM_MODE testing)
+ConfigureTest(STREAM_IO_PARQUET_TEST streams/io/parquet_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_JOIN_TEST streams/join_test.cpp STREAM_MODE testing)
-ConfigureTest(STREAM_JSONIO_TEST streams/io/json_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_LABELING_BINS_TEST streams/labeling_bins_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_LISTS_TEST streams/lists_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_MERGE_TEST streams/merge_test.cpp STREAM_MODE testing)
-ConfigureTest(STREAM_MULTIBYTE_SPLIT_TEST streams/io/multibyte_split_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_NULL_MASK_TEST streams/null_mask_test.cpp STREAM_MODE testing)
-ConfigureTest(STREAM_ORCIO_TEST streams/io/orc_test.cpp STREAM_MODE testing)
-ConfigureTest(STREAM_PARQUETIO_TEST streams/io/parquet_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_PARTITIONING_TEST streams/partitioning_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_POOL_TEST streams/pool_test.cu STREAM_MODE testing)
 ConfigureTest(STREAM_QUANTILE_TEST streams/quantile_test.cpp STREAM_MODE testing)

cpp/tests/large_strings/large_strings_fixture.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -127,5 +127,13 @@ int main(int argc, char** argv)
   // create object to automatically be destroyed at the end of main()
   auto lsd = cudf::test::StringsLargeTest::get_ls_data();
 
+  if (std::getenv("GTEST_CUDF_MEMORY_PEAK")) {
+    auto mr = rmm::mr::statistics_resource_adaptor<rmm::mr::device_memory_resource>(
+      cudf::get_current_device_resource());
+    cudf::set_current_device_resource(&mr);
+    auto rc = RUN_ALL_TESTS();
+    std::cout << "Peak memory usage " << mr.get_bytes_counter().peak << " bytes" << std::endl;
+    return rc;
+  }
   return RUN_ALL_TESTS();
 }

cpp/tests/streams/binaryop_test.cpp

@@ -17,6 +17,7 @@
 #include <cudf_test/base_fixture.hpp>
 #include <cudf_test/column_wrapper.hpp>
 #include <cudf_test/default_stream.hpp>
+#include <cudf_test/testing_main.hpp>
 
 #include <cudf/binaryop.hpp>
 #include <cudf/jit/runtime_support.hpp>
@@ -122,5 +123,8 @@ TEST_F(BinaryopPTXTest, ColumnColumnPTX)
 
   cudf::binary_operation(
     lhs, rhs, ptx, cudf::data_type(cudf::type_to_id<int32_t>()), cudf::test::get_default_stream());
-  cudf::binary_operation(lhs, rhs, ptx, cudf::data_type(cudf::type_to_id<int64_t>()));
+  cudf::binary_operation(
+    lhs, rhs, ptx, cudf::data_type(cudf::type_to_id<int64_t>()), cudf::test::get_default_stream());
 }
+
+CUDF_TEST_PROGRAM_MAIN()

cpp/tests/streams/concatenate_test.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #include <cudf_test/base_fixture.hpp>
 #include <cudf_test/column_wrapper.hpp>
 #include <cudf_test/default_stream.hpp>
+#include <cudf_test/testing_main.hpp>
 
 #include <cudf/concatenate.hpp>
 
@@ -49,3 +50,5 @@ TEST_F(ConcatenateTest, Masks)
   std::vector<cudf::column_view> views{input1, input2};
   auto result = cudf::concatenate_masks(views, cudf::test::get_default_stream());
 }
+
+CUDF_TEST_PROGRAM_MAIN()

cpp/tests/streams/copying_test.cpp

@@ -18,6 +18,7 @@
 #include <cudf_test/column_wrapper.hpp>
 #include <cudf_test/default_stream.hpp>
 #include <cudf_test/iterator_utilities.hpp>
+#include <cudf_test/testing_main.hpp>
 #include <cudf_test/type_lists.hpp>
 
 #include <cudf/contiguous_split.hpp>
@@ -360,3 +361,5 @@ TEST_F(CopyingTest, ContiguousSplit)
   cudf::table_view tbl({col, col2});
   auto result = cudf::contiguous_split(tbl, splits, cudf::test::get_default_stream());
 }
+
+CUDF_TEST_PROGRAM_MAIN()

cpp/tests/streams/datetime_test.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #include <cudf_test/base_fixture.hpp>
 #include <cudf_test/column_wrapper.hpp>
 #include <cudf_test/default_stream.hpp>
+#include <cudf_test/testing_main.hpp>
 
 #include <cudf/datetime.hpp>
 #include <cudf/scalar/scalar_factories.hpp>
@@ -147,3 +148,5 @@ TEST_F(DatetimeTest, RoundDatetimes)
   cudf::datetime::round_datetimes(
     timestamps, cudf::datetime::rounding_frequency::HOUR, cudf::test::get_default_stream());
 }
+
+CUDF_TEST_PROGRAM_MAIN()

cpp/tests/streams/dictionary_test.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #include <cudf_test/base_fixture.hpp>
 #include <cudf_test/column_wrapper.hpp>
 #include <cudf_test/default_stream.hpp>
+#include <cudf_test/testing_main.hpp>
 
 #include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/dictionary/dictionary_factories.hpp>
@@ -149,3 +150,5 @@ TEST_F(DictionaryTest, MatchDictionaries)
   cudf::test::fixed_width_column_wrapper<int> keys_col({2, 6});
   cudf::dictionary::match_dictionaries(dicts, cudf::test::get_default_stream());
 }
+
+CUDF_TEST_PROGRAM_MAIN()

cpp/tests/streams/filling_test.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #include <cudf_test/base_fixture.hpp>
 #include <cudf_test/column_wrapper.hpp>
 #include <cudf_test/default_stream.hpp>
+#include <cudf_test/testing_main.hpp>
 
 #include <cudf/column/column_view.hpp>
 #include <cudf/filling.hpp>
@@ -74,3 +75,5 @@ TEST_F(FillingTest, CalendricalMonthSequence)
 
   cudf::calendrical_month_sequence(10, init, 2, cudf::test::get_default_stream());
 }
+
+CUDF_TEST_PROGRAM_MAIN()