[ML] Periodically report pytorch_inference RSS

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"};
 

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);

bin/pytorch_inference/Main.cc

@@ -36,13 +36,24 @@
 #include <torch/csrc/api/include/torch/types.h>
 #include <torch/script.h>
 
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
 #include <cstdint>
 #include <exception>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <string>
+#include <thread>
 
 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<std::mutex> 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<std::mutex> 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();

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
 

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());
 }
 

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

docs/changelog/3160.yaml

@@ -0,0 +1,5 @@
+area: Machine Learning
+issues: []
+pr: 3160
+summary: Periodically report `pytorch_inference` RSS
+type: enhancement