diff --git a/bin/pytorch_inference/CResultWriter.cc b/bin/pytorch_inference/CResultWriter.cc index 34389dad4..74cdc200a 100644 --- a/bin/pytorch_inference/CResultWriter.cc +++ b/bin/pytorch_inference/CResultWriter.cc @@ -30,7 +30,9 @@ const std::string CResultWriter::ACK{"ack"}; const std::string CResultWriter::ACKNOWLEDGED{"acknowledged"}; const std::string CResultWriter::NUM_ALLOCATIONS{"num_allocations"}; const std::string CResultWriter::NUM_THREADS_PER_ALLOCATION{"num_threads_per_allocation"}; -const std::string CResultWriter::PROCESS_STATS{"process_stats"}; +// This key must match the field name the Elasticsearch PyTorchResult parser expects +// for process stats (see InferenceProcessStats / PyTorchResult on the Java side). +const std::string CResultWriter::PROCESS_STATS{"stats"}; const std::string CResultWriter::MEMORY_RESIDENT_SET_SIZE{"memory_rss"}; const std::string CResultWriter::MEMORY_MAX_RESIDENT_SET_SIZE{"memory_max_rss"}; diff --git a/bin/pytorch_inference/CResultWriter.h b/bin/pytorch_inference/CResultWriter.h index 8d809dc9d..fba77f221 100644 --- a/bin/pytorch_inference/CResultWriter.h +++ b/bin/pytorch_inference/CResultWriter.h @@ -73,6 +73,12 @@ class CResultWriter : public TStringBufWriter { void writeSimpleAck(const std::string_view& requestId); //! Write memory usage information to the output stream. + //! + //! Emits both the current resident set size (\c memory_rss) and the peak + //! resident set size (\c memory_max_rss, the OS high-water mark). The peak is + //! transmitted explicitly rather than derived on the Elasticsearch side from + //! the stream of samples, so transient spikes between reports are not lost; + //! this is the memory signal used to keep model assignment OOM-safe. void writeProcessStats(const std::string_view& requestId, const std::size_t residentSetSize, const std::size_t maxResidentSetSize); diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index 5b5fcca9f..1bfd39e3f 100644 --- a/bin/pytorch_inference/Main.cc +++ b/bin/pytorch_inference/Main.cc @@ -36,13 +36,24 @@ #include #include +#include +#include +#include #include #include #include +#include #include #include +#include namespace { +//! How often the periodic memory reporter emits the process resident set size. +//! Elasticsearch aggregates these samples over a longer window, so a short +//! interval gives several samples per window at negligible cost (a single +//! resident-set-size read). +constexpr std::chrono::seconds MEMORY_REPORT_INTERVAL{10}; + void verifySafeModel(const torch::jit::script::Module& module_) { try { auto result = ml::torch::CModelGraphValidator::validate(module_); @@ -377,6 +388,33 @@ int main(int argc, char** argv) { LOG_DEBUG(<< "Using a single allocation"); } + // Periodically report the resident set size so Elasticsearch can track the + // process's actual (OS-reported) memory use and bound model assignment and + // adaptive scaling by real memory rather than an a priori estimate. The + // command loop below blocks on input (a getline in CCommandParser::ioLoop), + // so the report is emitted from a dedicated timer thread. The concurrent + // line writer is safe to use from this thread alongside the inference-result + // threads. Shutdown is prompt: the condition variable is signalled the + // moment the command loop returns. + std::atomic_bool stopMemoryReporter{false}; + std::mutex memoryReporterMutex; + std::condition_variable memoryReporterCondition; + std::thread memoryReporterThread{[&] { + std::unique_lock lock{memoryReporterMutex}; + while (stopMemoryReporter.load() == false) { + memoryReporterCondition.wait_for(lock, MEMORY_REPORT_INTERVAL, [&] { + return stopMemoryReporter.load(); + }); + if (stopMemoryReporter.load()) { + break; + } + resultWriter.writeProcessStats( + ml::torch::CCommandParser::RESERVED_REQUEST_ID, + ml::core::CProcessStats::residentSetSize(), + ml::core::CProcessStats::maxResidentSetSize()); + } + }}; + commandParser.ioLoop( [&module_, &resultWriter](ml::torch::CCommandParser::CRequestCacheInterface& cache, ml::torch::CCommandParser::SRequest request) -> bool { @@ -391,6 +429,17 @@ int main(int argc, char** argv) { resultWriter.writeError(requestId, message); }); + // Stop the periodic memory reporter before tearing down the rest of the + // process so it cannot write to a closing output stream. + { + std::lock_guard lock{memoryReporterMutex}; + stopMemoryReporter.store(true); + } + memoryReporterCondition.notify_all(); + if (memoryReporterThread.joinable()) { + memoryReporterThread.join(); + } + // Stopping the executor forces this to block until all work is done if (useImmediateExecutor == false) { ml::core::stopDefaultAsyncExecutor(); diff --git a/bin/pytorch_inference/evaluate.py b/bin/pytorch_inference/evaluate.py index 7e60ec13a..89ceb6e5b 100644 --- a/bin/pytorch_inference/evaluate.py +++ b/bin/pytorch_inference/evaluate.py @@ -440,15 +440,16 @@ def memory_usage(args): # first get memory request request_sizes.insert(0, 0) - print(f"num items in request, memory_max_rss, inference time (ms)") + print(f"num items in request, memory_rss, memory_max_rss, inference time (ms)") for result in result_docs: if 'result' in result: inference_count = inference_count +1 last_time = result['time_ms'] continue - if 'process_stats' in result: - print(f"{request_sizes[stats_count]},{result['process_stats']['memory_max_rss']},{last_time}") + if 'stats' in result: + stats = result['stats'] + print(f"{request_sizes[stats_count]},{stats['memory_rss']},{stats.get('memory_max_rss')},{last_time}") stats_count = stats_count +1 continue diff --git a/bin/pytorch_inference/unittest/CResultWriterTest.cc b/bin/pytorch_inference/unittest/CResultWriterTest.cc index 7803bbc39..43ebe71f2 100644 --- a/bin/pytorch_inference/unittest/CResultWriterTest.cc +++ b/bin/pytorch_inference/unittest/CResultWriterTest.cc @@ -51,10 +51,10 @@ BOOST_AUTO_TEST_CASE(testWriteProcessStats) { std::ostringstream output; { ml::torch::CResultWriter resultWriter{output}; - resultWriter.writeProcessStats("req3", 42, 54); + resultWriter.writeProcessStats("req3", 42, 64); } - BOOST_REQUIRE_EQUAL("[{\"request_id\":\"req3\",\"process_stats\":" - "{\"memory_rss\":42,\"memory_max_rss\":54}}\n]", + BOOST_REQUIRE_EQUAL("[{\"request_id\":\"req3\",\"stats\":" + "{\"memory_rss\":42,\"memory_max_rss\":64}}\n]", output.str()); } diff --git a/docs/CHANGELOG.asciidoc b/docs/CHANGELOG.asciidoc index e66c892e4..ae47645fb 100644 --- a/docs/CHANGELOG.asciidoc +++ b/docs/CHANGELOG.asciidoc @@ -28,6 +28,12 @@ //=== Regressions +== {es} version 9.6.0 + +=== Enhancements + +* Periodically report the pytorch_inference process resident set size, including the peak (OS high-water mark), so Elasticsearch can track real native memory use per trained model deployment. (See {ml-pull}3160[#3160].) + == {es} version 9.4.0 === Bug Fixes diff --git a/docs/changelog/3160.yaml b/docs/changelog/3160.yaml new file mode 100644 index 000000000..4aacfcf53 --- /dev/null +++ b/docs/changelog/3160.yaml @@ -0,0 +1,5 @@ +area: Machine Learning +issues: [] +pr: 3160 +summary: Periodically report `pytorch_inference` RSS +type: enhancement