Introduce a mode in which data_frame_analyzer binary only outputs memory estimations but does not perform any analysis

przemekwitek · przemekwitek · commit 462023dd0f6f · 2019-08-06T13:07:01.000+02:00
diff --git a/bin/data_frame_analyzer/CCmdLineParser.cc b/bin/data_frame_analyzer/CCmdLineParser.cc
@@ -20,6 +20,7 @@ const std::string CCmdLineParser::DESCRIPTION = "Usage: data_frame_analyzer [opt
 bool CCmdLineParser::parse(int argc,
                            const char* const* argv,
                            std::string& configFile,
+                           bool& memoryUsageEstimationOnly,
                            std::string& logProperties,
                            std::string& logPipe,
                            bool& lengthEncodedInput,
@@ -35,6 +36,7 @@ bool CCmdLineParser::parse(int argc,
             ("version", "Display version information and exit")
             ("config", boost::program_options::value<std::string>(),
                     "The configuration file")
+            ("memoryUsageEstimationOnly", "Whether to perform memory usage estimation only")
             ("logProperties", boost::program_options::value<std::string>(),
                     "Optional logger properties file")
             ("logPipe", boost::program_options::value<std::string>(),
@@ -66,6 +68,9 @@ bool CCmdLineParser::parse(int argc,
         if (vm.count("config") > 0) {
             configFile = vm["config"].as<std::string>();
         }
+        if (vm.count("memoryUsageEstimationOnly") > 0) {
+            memoryUsageEstimationOnly = true;
+        }
         if (vm.count("logProperties") > 0) {
             logProperties = vm["logProperties"].as<std::string>();
         }
diff --git a/bin/data_frame_analyzer/CCmdLineParser.h b/bin/data_frame_analyzer/CCmdLineParser.h
@@ -27,6 +27,7 @@ class CCmdLineParser {
     static bool parse(int argc,
                       const char* const* argv,
                       std::string& configFile,
+                      bool& memoryUsageEstimationOnly,
                       std::string& logProperties,
                       std::string& logPipe,
                       bool& lengthEncodedInput,
diff --git a/bin/data_frame_analyzer/Main.cc b/bin/data_frame_analyzer/Main.cc
@@ -30,6 +30,7 @@
 #include <api/CDataFrameOutliersRunner.h>
 #include <api/CIoManager.h>
 #include <api/CLengthEncodedInputParser.h>
+#include <api/CMemoryUsageEstimationResultJsonWriter.h>
 
 #include "CCmdLineParser.h"
 
@@ -86,6 +87,7 @@ int main(int argc, char** argv) {
 
     // Read command line options
     std::string configFile;
+    bool memoryUsageEstimationOnly(false);
     std::string logProperties;
     std::string logPipe;
     bool lengthEncodedInput(false);
@@ -94,7 +96,7 @@ int main(int argc, char** argv) {
     std::string outputFileName;
     bool isOutputFileNamedPipe(false);
     if (ml::data_frame_analyzer::CCmdLineParser::parse(
-            argc, argv, configFile, logProperties, logPipe, lengthEncodedInput, inputFileName,
+            argc, argv, configFile, memoryUsageEstimationOnly, logProperties, logPipe, lengthEncodedInput, inputFileName,
             isInputFileNamedPipe, outputFileName, isOutputFileNamedPipe) == false) {
         return EXIT_FAILURE;
     }
@@ -127,13 +129,6 @@ int main(int argc, char** argv) {
     }
 
     using TInputParserUPtr = std::unique_ptr<ml::api::CInputParser>;
-    auto inputParser{[lengthEncodedInput, &ioMgr]() -> TInputParserUPtr {
-        if (lengthEncodedInput) {
-            return std::make_unique<ml::api::CLengthEncodedInputParser>(ioMgr.inputStream());
-        }
-        return std::make_unique<ml::api::CCsvInputParser>(
-            ioMgr.inputStream(), ml::api::CCsvInputParser::COMMA);
-    }()};
 
     std::string analysisSpecificationJson;
     bool couldReadConfigFile;
@@ -145,6 +140,17 @@ int main(int argc, char** argv) {
 
     auto analysisSpecification =
         std::make_unique<ml::api::CDataFrameAnalysisSpecification>(analysisSpecificationJson);
+
+    if (memoryUsageEstimationOnly) {
+        const auto result = analysisSpecification->estimateMemoryUsage();
+        auto outStream = [&ioMgr]() {
+            return std::make_unique<ml::core::CJsonOutputStreamWrapper>(ioMgr.outputStream());
+        }();
+        ml::api::CMemoryUsageEstimationResultJsonWriter jsonWriter(*outStream);
+        jsonWriter.write(result);
+        return EXIT_SUCCESS;
+    }
+
     if (analysisSpecification->numberThreads() > 1) {
         ml::core::startDefaultAsyncExecutor(analysisSpecification->numberThreads());
     }
@@ -156,6 +162,13 @@ int main(int argc, char** argv) {
 
     CCleanUpOnExit::add(dataFrameAnalyzer.dataFrameDirectory());
 
+    auto inputParser{[lengthEncodedInput, &ioMgr]() -> TInputParserUPtr {
+        if (lengthEncodedInput) {
+            return std::make_unique<ml::api::CLengthEncodedInputParser>(ioMgr.inputStream());
+        }
+        return std::make_unique<ml::api::CCsvInputParser>(
+            ioMgr.inputStream(), ml::api::CCsvInputParser::COMMA);
+    }()};
     if (inputParser->readStreamIntoVecs(
             [&dataFrameAnalyzer](const auto& fieldNames, const auto& fieldValues) {
                 return dataFrameAnalyzer.handleRecord(fieldNames, fieldValues);
diff --git a/include/api/CDataFrameAnalysisRunner.h b/include/api/CDataFrameAnalysisRunner.h
@@ -31,6 +31,7 @@ class CRowRef;
 }
 namespace api {
 class CDataFrameAnalysisSpecification;
+struct SMemoryUsageEstimationResult;
 
 //! \brief Hierarchy for running a specific core::CDataFrame analyses.
 //!
@@ -75,6 +76,11 @@ class API_EXPORT CDataFrameAnalysisRunner {
     //! number of rows per subset.
     void computeAndSaveExecutionStrategy();
 
+    //! Estimates memory usage in two cases: one partition (the whole data frame
+    //! fits in main memory) and maximum tolerable number of partitions (only
+    //! one partition needs to be loaded to main memory.
+    SMemoryUsageEstimationResult estimateMemoryUsage() const;
+
     //! Check if the data frame for this analysis should use in or out of core
     //! storage.
     bool storeDataFrameInMainMemory() const;
diff --git a/include/api/CDataFrameAnalysisSpecification.h b/include/api/CDataFrameAnalysisSpecification.h
@@ -151,6 +151,11 @@ class API_EXPORT CDataFrameAnalysisSpecification {
     //! calling thread until the runner has finished.
     CDataFrameAnalysisRunner* run(core::CDataFrame& frame) const;
 
+    //! Estimates memory usage in two cases: one partition (the whole data frame
+    //! fits in main memory) and maximum tolerable number of partitions (only
+    //! one partition needs to be loaded to main memory.
+    SMemoryUsageEstimationResult estimateMemoryUsage() const;
+
 private:
     void initializeRunner(const rapidjson::Value& jsonAnalysis);
 
diff --git a/include/api/CMemoryUsageEstimationResultJsonWriter.h b/include/api/CMemoryUsageEstimationResultJsonWriter.h
@@ -0,0 +1,50 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+#ifndef INCLUDED_ml_api_CMemoryUsageEstimationResultJsonWriter_h
+#define INCLUDED_ml_api_CMemoryUsageEstimationResultJsonWriter_h
+
+#include <core/CJsonOutputStreamWrapper.h>
+#include <core/CNonCopyable.h>
+#include <core/CRapidJsonConcurrentLineWriter.h>
+
+#include <api/ImportExport.h>
+
+#include <string>
+
+namespace ml {
+namespace api {
+
+//! Structure to store the memory usage estimation result
+struct API_EXPORT SMemoryUsageEstimationResult {
+    SMemoryUsageEstimationResult(size_t memoryUsageWithOnePartition, size_t memoryUsageWithMaxPartitions);
+    SMemoryUsageEstimationResult(const SMemoryUsageEstimationResult& result);
+
+    const size_t s_MemoryUsageWithOnePartition;
+    const size_t s_MemoryUsageWithMaxPartitions;
+};
+
+//! \brief
+//! Write memory usage estimation result in JSON format
+//!
+//! DESCRIPTION:\n
+//! Outputs the memory usage estimation result.
+//!
+class API_EXPORT CMemoryUsageEstimationResultJsonWriter : private core::CNonCopyable {
+public:
+    //! Constructor that causes output to be written to the specified wrapped stream
+    CMemoryUsageEstimationResultJsonWriter(core::CJsonOutputStreamWrapper& strmOut);
+
+    //! Writes the given memory usage estimation result in JSON format.
+    void write(const SMemoryUsageEstimationResult& result);
+
+private:
+    //! JSON line writer
+    core::CRapidJsonConcurrentLineWriter m_Writer;
+};
+}
+}
+
+#endif // INCLUDED_ml_api_CMemoryUsageEstimationResultJsonWriter_h
diff --git a/lib/api/CDataFrameAnalysisRunner.cc b/lib/api/CDataFrameAnalysisRunner.cc
@@ -11,6 +11,7 @@
 #include <core/CScopedFastLock.h>
 
 #include <api/CDataFrameAnalysisSpecification.h>
+#include <api/CMemoryUsageEstimationResultJsonWriter.h>
 
 #include <boost/iterator/counting_iterator.hpp>
 
@@ -36,6 +37,23 @@ CDataFrameAnalysisRunner::~CDataFrameAnalysisRunner() {
     this->waitToFinish();
 }
 
+SMemoryUsageEstimationResult CDataFrameAnalysisRunner::estimateMemoryUsage() const {
+    std::size_t numberRows{m_Spec.numberRows()};
+    std::size_t numberColumns{m_Spec.numberColumns() + this->numberExtraColumns()};
+    std::size_t maximumNumberPartitions{
+        static_cast<std::size_t>(std::sqrt(static_cast<double>(numberRows)) + 0.5)};
+    if (maximumNumberPartitions == 0) {
+        return SMemoryUsageEstimationResult(0, 0);
+    }
+    std::size_t memoryUsageWithOnePartition{
+        this->estimateMemoryUsage(numberRows, numberRows, numberColumns)};
+    std::size_t memoryUsageWithMaxPartitions{
+        this->estimateMemoryUsage(numberRows, numberRows / maximumNumberPartitions, numberColumns)};
+    return SMemoryUsageEstimationResult(
+        std::ceil(roundMb(memoryUsageWithOnePartition)),
+        std::ceil(roundMb(memoryUsageWithMaxPartitions)));
+}
+
 void CDataFrameAnalysisRunner::computeAndSaveExecutionStrategy() {
 
     std::size_t numberRows{m_Spec.numberRows()};
diff --git a/lib/api/CDataFrameAnalysisSpecification.cc b/lib/api/CDataFrameAnalysisSpecification.cc
@@ -12,6 +12,7 @@
 #include <api/CDataFrameAnalysisConfigReader.h>
 #include <api/CDataFrameBoostedTreeRunner.h>
 #include <api/CDataFrameOutliersRunner.h>
+#include <api/CMemoryUsageEstimationResultJsonWriter.h>
 
 #include <rapidjson/document.h>
 #include <rapidjson/ostreamwrapper.h>
@@ -183,6 +184,11 @@ CDataFrameAnalysisRunner* CDataFrameAnalysisSpecification::run(core::CDataFrame&
     return nullptr;
 }
 
+SMemoryUsageEstimationResult CDataFrameAnalysisSpecification::estimateMemoryUsage() const {
+    // TODO: What to do if m_Runner is nullptr?
+    return m_Runner->estimateMemoryUsage();
+}
+
 void CDataFrameAnalysisSpecification::initializeRunner(const rapidjson::Value& jsonAnalysis) {
     // We pass of the interpretation of the parameters object to the appropriate
     // analysis runner.
diff --git a/lib/api/CMemoryUsageEstimationResultJsonWriter.cc b/lib/api/CMemoryUsageEstimationResultJsonWriter.cc
@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+#include <api/CMemoryUsageEstimationResultJsonWriter.h>
+
+namespace ml {
+namespace api {
+namespace {
+
+// JSON field names
+const std::string MEMORY_USAGE_WITH_ONE_PARTITION("memory_usage_with_one_partition");
+const std::string MEMORY_USAGE_WITH_MAX_PARTITIONS("memory_usage_with_max_partitions");
+}
+
+SMemoryUsageEstimationResult::SMemoryUsageEstimationResult(
+    std::size_t memoryUsageWithOnePartition, std::size_t memoryUsageWithMaxPartitions):
+s_MemoryUsageWithOnePartition(memoryUsageWithOnePartition), s_MemoryUsageWithMaxPartitions(memoryUsageWithMaxPartitions) {}
+
+SMemoryUsageEstimationResult::SMemoryUsageEstimationResult(
+    const SMemoryUsageEstimationResult& result):
+SMemoryUsageEstimationResult(result.s_MemoryUsageWithOnePartition, result.s_MemoryUsageWithMaxPartitions) {}
+
+CMemoryUsageEstimationResultJsonWriter::CMemoryUsageEstimationResultJsonWriter(core::CJsonOutputStreamWrapper& strmOut) : m_Writer(strmOut) {
+    // Don't write any output in the constructor because, the way things work at
+    // the moment, the output stream might be redirected after construction
+}
+
+void CMemoryUsageEstimationResultJsonWriter::write(const SMemoryUsageEstimationResult& result) {
+    m_Writer.StartObject();
+    m_Writer.Key(MEMORY_USAGE_WITH_ONE_PARTITION);
+    m_Writer.Uint64(result.s_MemoryUsageWithOnePartition);
+    m_Writer.Key(MEMORY_USAGE_WITH_MAX_PARTITIONS);
+    m_Writer.Uint64(result.s_MemoryUsageWithMaxPartitions);
+    m_Writer.EndObject();
+    m_Writer.flush();
+}
+}
+}
diff --git a/lib/api/Makefile.first b/lib/api/Makefile.first
@@ -43,6 +43,7 @@ CInputParser.cc \
 CIoManager.cc \
 CJsonOutputWriter.cc \
 CLengthEncodedInputParser.cc \
+CMemoryUsageEstimationResultJsonWriter.cc \
 CModelPlotDataJsonWriter.cc \
 CModelSizeStatsJsonWriter.cc \
 CModelSnapshotJsonWriter.cc \
diff --git a/lib/api/unittest/CDataFrameAnalysisRunnerTest.cc b/lib/api/unittest/CDataFrameAnalysisRunnerTest.cc
@@ -13,6 +13,7 @@
 #include <api/CDataFrameAnalysisSpecification.h>
 #include <api/CDataFrameAnalysisSpecificationJsonWriter.h>
 #include <api/CDataFrameOutliersRunner.h>
+#include <api/CMemoryUsageEstimationResultJsonWriter.h>
 
 #include <test/CTestTmpDir.h>
 
@@ -121,6 +122,67 @@ void CDataFrameAnalysisRunnerTest::testComputeAndSaveExecutionStrategyDiskUsageF
     }
 }
 
+void CDataFrameAnalysisRunnerTest::testEstimateMemoryUsage() {
+
+    std::vector<std::string> errors;
+    std::mutex errorsMutex;
+    auto errorHandler = [&errors, &errorsMutex](std::string error) {
+        std::lock_guard<std::mutex> lock{errorsMutex};
+        errors.push_back(error);
+    };
+
+    core::CLogger::CScopeSetFatalErrorHandler scope{errorHandler};
+    api::CDataFrameOutliersRunnerFactory factory;
+
+    // Test estimation for empty data frame
+    {
+        errors.clear();
+        std::string jsonSpec{api::CDataFrameAnalysisSpecificationJsonWriter::jsonString(
+            0, 5, 100000000, 1, {}, true, test::CTestTmpDir::tmpDir(), "", "outlier_detection", "")};
+        api::CDataFrameAnalysisSpecification spec{jsonSpec};
+
+        // Check memory estimation result
+        api::SMemoryUsageEstimationResult result = spec.estimateMemoryUsage();
+        CPPUNIT_ASSERT_EQUAL(0, static_cast<int>(result.s_MemoryUsageWithOnePartition));
+        CPPUNIT_ASSERT_EQUAL(0, static_cast<int>(result.s_MemoryUsageWithMaxPartitions));
+
+        // no error should be registered
+        CPPUNIT_ASSERT_EQUAL(1, static_cast<int>(errors.size()));
+    }
+
+    // Test estimation for data frame with 1 row
+    {
+        errors.clear();
+        std::string jsonSpec{api::CDataFrameAnalysisSpecificationJsonWriter::jsonString(
+            1, 5, 100000000, 1, {}, true, test::CTestTmpDir::tmpDir(), "", "outlier_detection", "")};
+        api::CDataFrameAnalysisSpecification spec{jsonSpec};
+
+        // Check memory estimation result
+        api::SMemoryUsageEstimationResult result = spec.estimateMemoryUsage();
+        CPPUNIT_ASSERT_EQUAL(6050, static_cast<int>(result.s_MemoryUsageWithOnePartition));
+        CPPUNIT_ASSERT_EQUAL(6050, static_cast<int>(result.s_MemoryUsageWithMaxPartitions));
+
+        // no error should be registered
+        CPPUNIT_ASSERT_EQUAL(0, static_cast<int>(errors.size()));
+    }
+
+    // Test estimation for data frame with 4 rows
+    {
+        errors.clear();
+        std::string jsonSpec{api::CDataFrameAnalysisSpecificationJsonWriter::jsonString(
+            4, 5, 100000000, 1, {}, true, test::CTestTmpDir::tmpDir(), "", "outlier_detection", "")};
+        api::CDataFrameAnalysisSpecification spec{jsonSpec};
+
+        // Check memory estimation result
+        api::SMemoryUsageEstimationResult result = spec.estimateMemoryUsage();
+        CPPUNIT_ASSERT_EQUAL(9104, static_cast<int>(result.s_MemoryUsageWithOnePartition));
+        CPPUNIT_ASSERT_EQUAL(8528, static_cast<int>(result.s_MemoryUsageWithMaxPartitions));
+
+        // no error should be registered
+        CPPUNIT_ASSERT_EQUAL(0, static_cast<int>(errors.size()));
+    }
+}
+
 CppUnit::Test* CDataFrameAnalysisRunnerTest::suite() {
     CppUnit::TestSuite* suiteOfTests = new CppUnit::TestSuite("CDataFrameAnalysisRunnerTest");
 
@@ -130,6 +192,9 @@ CppUnit::Test* CDataFrameAnalysisRunnerTest::suite() {
     suiteOfTests->addTest(new CppUnit::TestCaller<CDataFrameAnalysisRunnerTest>(
         "CDataFrameAnalysisRunnerTest::testComputeAndSaveExecutionStrategyDiskUsageFlag",
         &CDataFrameAnalysisRunnerTest::testComputeAndSaveExecutionStrategyDiskUsageFlag));
+    suiteOfTests->addTest(new CppUnit::TestCaller<CDataFrameAnalysisRunnerTest>(
+        "CDataFrameAnalysisRunnerTest::testEstimateMemoryUsage",
+        &CDataFrameAnalysisRunnerTest::testEstimateMemoryUsage));
 
     return suiteOfTests;
 }
diff --git a/lib/api/unittest/CDataFrameAnalysisRunnerTest.h b/lib/api/unittest/CDataFrameAnalysisRunnerTest.h
@@ -13,6 +13,7 @@ class CDataFrameAnalysisRunnerTest : public CppUnit::TestFixture {
 public:
     void testComputeExecutionStrategyForOutliers();
     void testComputeAndSaveExecutionStrategyDiskUsageFlag();
+    void testEstimateMemoryUsage();
 
     static CppUnit::Test* suite();
 
diff --git a/lib/api/unittest/CMemoryUsageEstimationResultJsonWriterTest.cc b/lib/api/unittest/CMemoryUsageEstimationResultJsonWriterTest.cc
diff --git a/lib/api/unittest/CMemoryUsageEstimationResultJsonWriterTest.h b/lib/api/unittest/CMemoryUsageEstimationResultJsonWriterTest.h
diff --git a/lib/api/unittest/Main.cc b/lib/api/unittest/Main.cc
diff --git a/lib/api/unittest/Makefile b/lib/api/unittest/Makefile
