[7.8][ML] Eagerly discard node statistics for leaves which we will never split (elastic#1148)

Backport elastic#1125.

Author: tveasey

docs/CHANGELOG.asciidoc

@@ -41,9 +41,11 @@
   (See {ml-pull}1113[#1113].)
 * Improve robustness of anomaly detection to bad input data. (See {ml-pull}1114[#1114].)
 * Adds new `num_matches` and `preferred_to_categories` fields to category output.
-  (See {ml-pull}1062[#1062])
+  (See {ml-pull}1062[#1062].)
 * Adds mean squared logarithmic error (MSLE) for regression. (See {ml-pull}1101[#1101].)
 * Improve robustness of anomaly detection to bad input data. (See {ml-pull}1114[#1114].)
+* Reduce peak memory usage and memory estimates for classification and regression.
+  (See {ml-pull}1125[#1125].)
 * Switched data frame analytics model memory estimates from kilobytes to megabytes.
   (See {ml-pull}1126[#1126], issue: {issue}54506[#54506].)
 * Added a {ml} native code build for Linux on AArch64. (See {ml-pull}1132[#1132] and

lib/api/unittest/CDataFrameAnalyzerTrainingTest.cc

@@ -330,8 +330,8 @@ BOOST_AUTO_TEST_CASE(testRunBoostedTreeRegressionTrainingMse) {
               << "ms");
 
     BOOST_TEST_REQUIRE(core::CProgramCounters::counter(
-                           counter_t::E_DFTPMEstimatedPeakMemoryUsage) < 6000000);
-    BOOST_TEST_REQUIRE(core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) < 1500000);
+                           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) <
         core::CProgramCounters::counter(counter_t::E_DFTPMEstimatedPeakMemoryUsage));
@@ -720,8 +720,8 @@ BOOST_AUTO_TEST_CASE(testRunBoostedTreeClassifierTraining) {
               << "ms");
 
     BOOST_TEST_REQUIRE(core::CProgramCounters::counter(
-                           counter_t::E_DFTPMEstimatedPeakMemoryUsage) < 6000000);
-    BOOST_TEST_REQUIRE(core::CProgramCounters::counter(counter_t::E_DFTPMPeakMemoryUsage) < 1500000);
+                           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) <
         core::CProgramCounters::counter(counter_t::E_DFTPMEstimatedPeakMemoryUsage));

lib/api/unittest/CDataFrameMockAnalysisRunner.h

@@ -13,13 +13,14 @@
 
 #include <test/CRandomNumbers.h>
 
-#include <functional>
+#include <cinttypes>
+#include <string>
 
 class CDataFrameMockAnalysisState final : public ml::api::CDataFrameAnalysisInstrumentation {
 public:
     CDataFrameMockAnalysisState(const std::string& jobId)
         : ml::api::CDataFrameAnalysisInstrumentation(jobId) {}
-    void writeAnalysisStats(std::int64_t /* timestamp */) override{};
+    void writeAnalysisStats(std::int64_t /* timestamp */) override {}
 
 protected:
     ml::counter_t::ECounterTypes memoryCounterType() override;

lib/maths/CBoostedTreeImpl.cc

@@ -4,7 +4,6 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-#include "maths/CBoostedTreeUtils.h"
 #include <maths/CBoostedTreeImpl.h>
 
 #include <core/CContainerPrinter.h>
@@ -20,13 +19,18 @@
 #include <maths/CBoostedTree.h>
 #include <maths/CBoostedTreeLeafNodeStatistics.h>
 #include <maths/CBoostedTreeLoss.h>
+#include <maths/CBoostedTreeUtils.h>
 #include <maths/CDataFrameAnalysisInstrumentationInterface.h>
 #include <maths/CDataFrameCategoryEncoder.h>
 #include <maths/CQuantileSketch.h>
 #include <maths/CSampling.h>
 #include <maths/CSetTools.h>
 #include <maths/CTreeShapFeatureImportance.h>
 
+#include <boost/circular_buffer.hpp>
+
+#include <algorithm>
+
 namespace ml {
 namespace maths {
 using namespace boosted_tree;
@@ -310,10 +314,14 @@ std::size_t CBoostedTreeImpl::estimateMemoryUsage(std::size_t numberRows,
     std::size_t foldRoundLossMemoryUsage{m_NumberFolds * m_NumberRounds *
                                          sizeof(TOptionalDouble)};
     std::size_t hyperparametersMemoryUsage{numberColumns * sizeof(double)};
+    // We only maintain statistics for leaves we know we may possibly split this
+    // halves the peak number of statistics we maintain.
     std::size_t leafNodeStatisticsMemoryUsage{
-        maximumNumberLeaves * CBoostedTreeLeafNodeStatistics::estimateMemoryUsage(
-                                  numberRows, maximumNumberFeatures, m_NumberSplitsPerFeature,
-                                  m_Loss->numberParameters())};
+        maximumNumberLeaves *
+        CBoostedTreeLeafNodeStatistics::estimateMemoryUsage(
+            numberRows, maximumNumberFeatures, m_NumberSplitsPerFeature,
+            m_Loss->numberParameters()) /
+        2};
     std::size_t dataTypeMemoryUsage{maximumNumberFeatures * sizeof(CDataFrameUtils::SDataType)};
     std::size_t featureSampleProbabilities{maximumNumberFeatures * sizeof(double)};
     std::size_t missingFeatureMaskMemoryUsage{
@@ -721,14 +729,13 @@ CBoostedTreeImpl::trainTree(core::CDataFrame& frame,
     LOG_TRACE(<< "Training one tree...");
 
     using TLeafNodeStatisticsPtr = CBoostedTreeLeafNodeStatistics::TPtr;
-    using TLeafNodeStatisticsPtrQueue =
-        std::priority_queue<TLeafNodeStatisticsPtr, std::vector<TLeafNodeStatisticsPtr>, COrderings::SLess>;
+    using TLeafNodeStatisticsPtrQueue = boost::circular_buffer<TLeafNodeStatisticsPtr>;
 
     TNodeVec tree(1);
     tree.reserve(2 * maximumTreeSize + 1);
 
-    TLeafNodeStatisticsPtrQueue leaves;
-    leaves.push(std::make_shared<CBoostedTreeLeafNodeStatistics>(
+    TLeafNodeStatisticsPtrQueue leaves(maximumTreeSize / 2 + 3);
+    leaves.push_back(std::make_shared<CBoostedTreeLeafNodeStatistics>(
         0 /*root*/, m_NumberInputColumns, m_Loss->numberParameters(),
         m_NumberThreads, frame, *m_Encoder, m_Regularization, candidateSplits,
         this->featureBag(), 0 /*depth*/, trainingRowMask));
@@ -752,10 +759,16 @@ CBoostedTreeImpl::trainTree(core::CDataFrame& frame,
 
     double totalGain{0.0};
 
+    COrderings::SLess less;
+
     for (std::size_t i = 0; i < maximumTreeSize; ++i) {
 
-        auto leaf = leaves.top();
-        leaves.pop();
+        if (leaves.empty()) {
+            break;
+        }
+
+        auto leaf = leaves.back();
+        leaves.pop_back();
 
         scopeMemoryUsage.remove(leaf);
 
@@ -764,7 +777,8 @@ CBoostedTreeImpl::trainTree(core::CDataFrame& frame,
         }
 
         totalGain += leaf->gain();
-        LOG_TRACE(<< "splitting " << leaf->id() << " total gain = " << totalGain);
+        LOG_TRACE(<< "splitting " << leaf->id() << " leaf gain = " << leaf->gain()
+                  << " total gain = " << totalGain);
 
         std::size_t splitFeature;
         double splitValue;
@@ -783,11 +797,26 @@ CBoostedTreeImpl::trainTree(core::CDataFrame& frame,
             leftChildId, rightChildId, m_NumberThreads, frame, *m_Encoder,
             m_Regularization, candidateSplits, this->featureBag(), tree[leaf->id()]);
 
-        scopeMemoryUsage.add(leftChild);
-        scopeMemoryUsage.add(rightChild);
+        if (less(rightChild, leftChild)) {
+            std::swap(leftChild, rightChild);
+        }
 
-        leaves.push(std::move(leftChild));
-        leaves.push(std::move(rightChild));
+        std::size_t n{leaves.size()};
+        if (leftChild->gain() >= MINIMUM_RELATIVE_GAIN_PER_SPLIT * totalGain) {
+            scopeMemoryUsage.add(leftChild);
+            leaves.push_back(std::move(leftChild));
+        }
+        if (rightChild->gain() >= MINIMUM_RELATIVE_GAIN_PER_SPLIT * totalGain) {
+            scopeMemoryUsage.add(rightChild);
+            leaves.push_back(std::move(rightChild));
+        }
+        std::inplace_merge(leaves.begin(), leaves.begin() + n, leaves.end(), less);
+
+        // Drop any leaves which can't possibly be split.
+        while (leaves.size() + i + 1 > maximumTreeSize) {
+            scopeMemoryUsage.remove(leaves.front());
+            leaves.pop_front();
+        }
     }
 
     tree.shrink_to_fit();