[7.8][ML] Fix memory usage estimation for vectors with a custom allocator (#1162)

tveasey · web-flow · commit def6d14af6c2 · 2020-04-24T10:35:51.000+01:00
Backport #1156.
diff --git a/include/core/CMemory.h b/include/core/CMemory.h
@@ -327,8 +327,8 @@ class CORE_EXPORT CMemory : private CNonInstantiatable {
     }
 
     //! Overload for std::vector.
-    template<typename T>
-    static std::size_t dynamicSize(const std::vector<T>& t) {
+    template<typename T, typename A>
+    static std::size_t dynamicSize(const std::vector<T, A>& t) {
         std::size_t mem = 0;
         if (!memory_detail::SDynamicSizeAlwaysZero<T>::value()) {
             for (auto i = t.begin(); i != t.end(); ++i) {
@@ -781,9 +781,9 @@ class CORE_EXPORT CMemoryDebug : private CNonInstantiatable {
     }
 
     //! Overload for std::vector.
-    template<typename T>
+    template<typename T, typename A>
     static void dynamicSize(const char* name,
-                            const std::vector<T>& t,
+                            const std::vector<T, A>& t,
                             const CMemoryUsage::TMemoryUsagePtr& mem) {
         std::string componentName(name);
 
diff --git a/lib/api/CDataFrameTrainBoostedTreeClassifierRunner.cc b/lib/api/CDataFrameTrainBoostedTreeClassifierRunner.cc
@@ -162,9 +162,9 @@ void CDataFrameTrainBoostedTreeClassifierRunner::writeOneRow(
     }
 
     if (featureImportance != nullptr) {
+        int numberClasses{static_cast<int>(classValues.size())};
         featureImportance->shap(
-            row, [&writer, &classValues](
-                     const maths::CTreeShapFeatureImportance::TSizeVec& indices,
+            row, [&](const maths::CTreeShapFeatureImportance::TSizeVec& indices,
                      const TStrVec& featureNames,
                      const maths::CTreeShapFeatureImportance::TVectorVec& shap) {
                 writer.Key(FEATURE_IMPORTANCE_FIELD_NAME);
@@ -178,8 +178,7 @@ void CDataFrameTrainBoostedTreeClassifierRunner::writeOneRow(
                             writer.Key(IMPORTANCE_FIELD_NAME);
                             writer.Double(shap[i](0));
                         } else {
-                            for (int j = 0;
-                                 j < shap[i].size() && j < classValues.size(); ++j) {
+                            for (int j = 0; j < shap[i].size() && j < numberClasses; ++j) {
                                 writer.Key(classValues[j]);
                                 writer.Double(shap[i](j));
                             }
diff --git a/lib/api/unittest/CDataFrameAnalyzerTrainingTest.cc b/lib/api/unittest/CDataFrameAnalyzerTrainingTest.cc
@@ -331,7 +331,7 @@ BOOST_AUTO_TEST_CASE(testRunBoostedTreeRegressionTrainingMse) {
 
     BOOST_TEST_REQUIRE(core::CProgramCounters::counter(
                            counter_t::E_DFTPMEstimatedPeakMemoryUsage) < 4500000);
-    BOOST_TEST_REQUIRE(core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) < 1600000);
+    BOOST_TEST_REQUIRE(core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) < 1800000);
     BOOST_TEST_REQUIRE(
         core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) <
         core::CProgramCounters::counter(counter_t::E_DFTPMEstimatedPeakMemoryUsage));
@@ -721,7 +721,7 @@ BOOST_AUTO_TEST_CASE(testRunBoostedTreeClassifierTraining) {
 
     BOOST_TEST_REQUIRE(core::CProgramCounters::counter(
                            counter_t::E_DFTPMEstimatedPeakMemoryUsage) < 4500000);
-    BOOST_TEST_REQUIRE(core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) < 1600000);
+    BOOST_TEST_REQUIRE(core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) < 1800000);
     BOOST_TEST_REQUIRE(
         core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) <
         core::CProgramCounters::counter(counter_t::E_DFTPMEstimatedPeakMemoryUsage));
diff --git a/lib/core/unittest/CDataFrameTest.cc b/lib/core/unittest/CDataFrameTest.cc
@@ -427,9 +427,10 @@ BOOST_FIXTURE_TEST_CASE(testMemoryUsage, CTestFixture) {
 
     // Memory usage should be less than:
     //   1) 1075 + 4 times the root directory length bytes for on disk, and
-    //   2) data size + doc ids size + 900 byte overhead in main memory.
+    //   2) data size + doc ids size + 900 byte overhead in main memory. Note that
+    //      we round up the number of columns to the next multiple of the alignment.
     std::size_t maximumMemory[]{1075 + 4 * rootDirectory.length(),
-                                rows * (cols + 1) * 4 + 900};
+                                rows * (4 * ((cols + 3) / 4) + 1) * 4 + 900};
 
     std::string type[]{"on disk", "main memory"};
     std::size_t t{0};
diff --git a/lib/core/unittest/CMemoryUsageTest.cc b/lib/core/unittest/CMemoryUsageTest.cc
@@ -4,6 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+#include <core/CAlignment.h>
 #include <core/CContainerPrinter.h>
 #include <core/CHashing.h>
 #include <core/CLogger.h>
@@ -167,8 +168,10 @@ class CBase {
 
     virtual std::size_t staticSize() const { return sizeof(*this); }
 
+    const std::uint64_t* fixed() const { return m_Fixed; } // suppress warning
+
 private:
-    uint64_t m_Fixed[5];
+    std::uint64_t m_Fixed[5];
     TIntVec m_Vec;
 };
 
@@ -193,8 +196,10 @@ class CDerived : public CBase {
 
     virtual std::size_t staticSize() const { return sizeof(*this); }
 
+    const std::uint64_t* fixed() const { return m_Fixed; } // suppress warning
+
 private:
-    uint64_t m_Fixed[50];
+    std::uint64_t m_Fixed[50];
     TStrVec m_Strings;
 };
 
@@ -1204,4 +1209,27 @@ BOOST_AUTO_TEST_CASE(testSmallVector) {
     BOOST_TEST_REQUIRE(extraMem > 0);
 }
 
+BOOST_AUTO_TEST_CASE(testAlignedVector) {
+    using TDoubleVec = std::vector<double>;
+    using TAlignedDoubleVec = std::vector<double, core::CAlignedAllocator<double>>;
+
+    TDoubleVec vector{10.0, 11.0, 12.0, 13.0, 14.0,
+                      15.0, 16.0, 17.0, 18.0, 19.0};
+    TAlignedDoubleVec alignedVector{10.0, 11.0, 12.0, 13.0, 14.0,
+                                    15.0, 16.0, 17.0, 18.0, 19.0};
+
+    LOG_DEBUG(<< "TDoubleVec usage = " << core::CMemory::dynamicSize(vector));
+    LOG_DEBUG(<< "TAlignedDoubleVec usage = " << core::CMemory::dynamicSize(alignedVector));
+    BOOST_REQUIRE_EQUAL(core::CMemory::dynamicSize(vector),
+                        core::CMemory::dynamicSize(alignedVector));
+
+    core::CMemoryUsage memoryUsage;
+    memoryUsage.setName("test", 0);
+    core::CMemoryDebug::dynamicSize("TAlignedDoubleVec", vector, memoryUsage.addChild());
+    std::ostringstream ss;
+    memoryUsage.print(ss);
+    LOG_DEBUG(<< "TAlignedDoubleVec usage debug = " << ss.str());
+    BOOST_REQUIRE_EQUAL(core::CMemory::dynamicSize(vector), memoryUsage.usage());
+}
+
 BOOST_AUTO_TEST_SUITE_END()
