From 2017e2c5c0d4e2aee278e8030c6bf962d4af45d2 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Mon, 17 Aug 2026 15:46:56 +1200 Subject: [PATCH 1/3] [ML] Periodically report pytorch_inference RSS Emit the pytorch_inference resident set size on a fixed 10s interval, independent of inference requests, reporting both the current RSS (memory_rss) and the OS peak (memory_max_rss). This lets Elasticsearch track real native memory use per trained model deployment and keep assignment and adaptive scaling OOM-safe rather than relying on an a priori estimate. The process-stats field is renamed to "stats" to match the Elasticsearch PyTorchResult parser. Relates elastic/ml-cpp#2885 --- bin/pytorch_inference/CResultWriter.cc | 4 +- bin/pytorch_inference/CResultWriter.h | 6 +++ bin/pytorch_inference/Main.cc | 49 +++++++++++++++++++ bin/pytorch_inference/evaluate.py | 7 +-- .../unittest/CResultWriterTest.cc | 6 +-- 5 files changed, 65 insertions(+), 7 deletions(-) diff --git a/bin/pytorch_inference/CResultWriter.cc b/bin/pytorch_inference/CResultWriter.cc index 34389dad44..74cdc200a7 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 8d809dc9df..fba77f2213 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 5b5fcca9f9..1bfd39e3f3 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 7e60ec13a1..89ceb6e5ba 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 7803bbc391..43ebe71f2a 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()); } From 72e6b2575e3c28a94fa2e70ff3b6b737a3f1702f Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Mon, 17 Aug 2026 15:48:09 +1200 Subject: [PATCH 2/3] [ML] Add changelog for periodic RSS reporting Relates elastic/ml-cpp#2885 --- docs/CHANGELOG.asciidoc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/CHANGELOG.asciidoc b/docs/CHANGELOG.asciidoc index e66c892e42..ae47645fb8 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 From 3d69f56289311d49c4f5153334955fc443b92059 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Fri, 21 Aug 2026 13:13:15 +1200 Subject: [PATCH 3/3] Update docs/changelog/3160.yaml --- docs/changelog/3160.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 docs/changelog/3160.yaml diff --git a/docs/changelog/3160.yaml b/docs/changelog/3160.yaml new file mode 100644 index 0000000000..4aacfcf536 --- /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