[Enhancement](pyudf) Support MemTracker in PythonUdf (#60655)

Author: linrrzqqq

be/src/common/config.cpp

@@ -1129,6 +1129,9 @@ DEFINE_String(python_venv_interpreter_paths, "");
 // max python processes in global shared pool, each version can have up to this many processes
 // 0 means use CPU core count as default, otherwise use the specified value
 DEFINE_mInt32(max_python_process_num, "0");
+// Memory limit in bytes for all Python UDF processes; warning is logged when exceeded
+// default is 10GB
+DEFINE_mInt64(python_udf_processes_memory_limit_bytes, "10737418240");
 
 // Set config randomly to check more issues in github workflow
 DEFINE_Bool(enable_fuzzy_mode, "false");

be/src/common/config.h

@@ -1190,6 +1190,8 @@ DECLARE_String(python_venv_root_path);
 DECLARE_String(python_venv_interpreter_paths);
 // max python processes in global shared pool, each version can have up to this many processes
 DECLARE_mInt32(max_python_process_num);
+// Memory limit in bytes for all Python UDF processes; warning is logged when exceeded
+DECLARE_mInt64(python_udf_processes_memory_limit_bytes);
 
 // Set config randomly to check more issues in github workflow
 DECLARE_Bool(enable_fuzzy_mode);

be/src/udf/python/python_server.cpp

@@ -26,6 +26,7 @@
 
 #include <boost/asio.hpp>
 #include <boost/process.hpp>
+#include <fstream>
 
 #include "common/config.h"
 #include "udf/python/python_udaf_client.h"
@@ -191,66 +192,70 @@ Status PythonServerManager::fork(const PythonVersion& version, ProcessPtr* proce
 void PythonServerManager::_start_health_check_thread() {
     if (_health_check_thread) return;
 
-    LOG(INFO) << "Starting Python process health check thread (interval: 60 seconds)";
+    LOG(INFO) << "Starting Python process health check thread (interval: 30 seconds)";
 
     _health_check_thread = std::make_unique<std::thread>([this]() {
         // Health check loop
         while (!_shutdown_flag.load(std::memory_order_acquire)) {
             // Wait for interval or shutdown signal
             {
                 std::unique_lock<std::mutex> lock(_health_check_mutex);
-                _health_check_cv.wait_for(lock, std::chrono::seconds(60), [this]() {
+                _health_check_cv.wait_for(lock, std::chrono::seconds(30), [this]() {
                     return _shutdown_flag.load(std::memory_order_acquire);
                 });
             }
 
             if (_shutdown_flag.load(std::memory_order_acquire)) break;
 
-            std::lock_guard<std::mutex> lock(_pools_mutex);
-
-            int total_checked = 0;
-            int total_dead = 0;
-            int total_recreated = 0;
-
-            for (auto& [version, pool] : _process_pools) {
-                for (size_t i = 0; i < pool.size(); ++i) {
-                    auto& process = pool[i];
-                    if (!process) continue;
-
-                    total_checked++;
-                    if (!process->is_alive()) {
-                        total_dead++;
-                        LOG(WARNING)
-                                << "Detected dead Python process (pid=" << process->get_child_pid()
-                                << ", version=" << version.to_string() << "), recreating...";
-
-                        ProcessPtr new_process;
-                        Status s = fork(version, &new_process);
-                        if (s.ok()) {
-                            pool[i] = std::move(new_process);
-                            total_recreated++;
-                            LOG(INFO) << "Successfully recreated Python process for version "
-                                      << version.to_string();
-                        } else {
-                            LOG(ERROR) << "Failed to recreate Python process for version "
-                                       << version.to_string() << ": " << s.to_string();
-                            pool.erase(pool.begin() + i);
-                            --i;
-                        }
-                    }
-                }
-            }
-
-            if (total_dead > 0) {
-                LOG(INFO) << "Health check completed: checked=" << total_checked
-                          << ", dead=" << total_dead << ", recreated=" << total_recreated;
-            }
+            _check_and_recreate_processes();
+            _refresh_memory_stats();
         }
 
         LOG(INFO) << "Python process health check thread exiting";
     });
 }
 
+void PythonServerManager::_check_and_recreate_processes() {
+    std::lock_guard<std::mutex> lock(_pools_mutex);
+
+    int total_checked = 0;
+    int total_dead = 0;
+    int total_recreated = 0;
+
+    for (auto& [version, pool] : _process_pools) {
+        for (size_t i = 0; i < pool.size(); ++i) {
+            auto& process = pool[i];
+            if (!process) continue;
+
+            total_checked++;
+            if (!process->is_alive()) {
+                total_dead++;
+                LOG(WARNING) << "Detected dead Python process (pid=" << process->get_child_pid()
+                             << ", version=" << version.to_string() << "), recreating...";
+
+                ProcessPtr new_process;
+                Status s = fork(version, &new_process);
+                if (s.ok()) {
+                    pool[i] = std::move(new_process);
+                    total_recreated++;
+                    LOG(INFO) << "Successfully recreated Python process for version "
+                              << version.to_string();
+                } else {
+                    LOG(ERROR) << "Failed to recreate Python process for version "
+                               << version.to_string() << ": " << s.to_string();
+                    pool.erase(pool.begin() + i);
+                    --i;
+                }
+            }
+        }
+    }
+
+    if (total_dead > 0) {
+        LOG(INFO) << "Health check completed: checked=" << total_checked << ", dead=" << total_dead
+                  << ", recreated=" << total_recreated;
+    }
+}
+
 void PythonServerManager::shutdown() {
     // Signal health check thread to stop
     _shutdown_flag.store(true, std::memory_order_release);
@@ -273,6 +278,61 @@ void PythonServerManager::shutdown() {
     _process_pools.clear();
 }
 
+Status PythonServerManager::_read_process_memory(pid_t pid, size_t* rss_bytes) {
+    // Read from /proc/{pid}/statm
+    // Format: size resident shared text lib data dt
+    std::string statm_path = fmt::format("/proc/{}/statm", pid);
+    std::ifstream statm_file(statm_path);
+
+    if (!statm_file.is_open()) {
+        return Status::InternalError("Cannot open {}", statm_path);
+    }
+
+    size_t size_pages = 0, rss_pages = 0;
+    // we only care about RSS, read and ignore the total size field
+    statm_file >> size_pages >> rss_pages;
+
+    if (statm_file.fail()) {
+        return Status::InternalError("Failed to read {}", statm_path);
+    }
+
+    // Convert pages to bytes
+    long page_size = sysconf(_SC_PAGESIZE);
+    *rss_bytes = rss_pages * page_size;
+
+    return Status::OK();
+}
+
+void PythonServerManager::_refresh_memory_stats() {
+    std::lock_guard<std::mutex> lock(_pools_mutex);
+
+    int64_t total_rss = 0;
+
+    for (const auto& [version, pool] : _process_pools) {
+        for (const auto& process : pool) {
+            if (!process || !process->is_alive()) continue;
+
+            size_t rss_bytes = 0;
+            Status s = _read_process_memory(process->get_child_pid(), &rss_bytes);
+
+            if (s.ok()) {
+                total_rss += rss_bytes;
+            } else [[unlikely]] {
+                LOG(WARNING) << "Failed to read memory info for Python process (pid="
+                             << process->get_child_pid() << "): " << s.to_string();
+            }
+        }
+    }
+    _mem_tracker.set_consumption(total_rss);
+    LOG(INFO) << _mem_tracker.log_usage();
+
+    if (config::python_udf_processes_memory_limit_bytes > 0 &&
+        total_rss > config::python_udf_processes_memory_limit_bytes) {
+        LOG(WARNING) << "Python UDF process memory usage exceeds limit: rss_bytes=" << total_rss
+                     << ", limit_bytes=" << config::python_udf_processes_memory_limit_bytes;
+    }
+}
+
 // Explicit template instantiation for UDF, UDAF and UDTF clients
 template Status PythonServerManager::get_client<PythonUDFClient>(
         const PythonUDFMeta& func_meta, const PythonVersion& version,

be/src/udf/python/python_server.h

@@ -23,11 +23,11 @@
 #include <thread>
 
 #include "common/status.h"
+#include "runtime/memory/mem_tracker.h"
 #include "udf/python/python_udf_meta.h"
 #include "udf/python/python_udf_runtime.h"
 
 namespace doris {
-
 class PythonServerManager {
 public:
     PythonServerManager() = default;
@@ -52,13 +52,37 @@ class PythonServerManager {
 
     void shutdown();
 
+#ifdef BE_TEST
+    // For unit testing only.
+    void check_and_recreate_processes_for_test() { _check_and_recreate_processes(); }
+
+    std::unordered_map<PythonVersion, std::vector<ProcessPtr>>& process_pools_for_test() {
+        return _process_pools;
+    }
+#endif
+
 private:
     /**
      * Start health check background thread (called once by ensure_pool_initialized)
-     * Thread periodically checks process health and recreates dead processes
+     * Thread periodically checks process health and refreshes memory stats
      */
     void _start_health_check_thread();
 
+    /**
+     * Check process health and recreate dead processes
+     */
+    void _check_and_recreate_processes();
+
+    /**
+     * Read resident set size (RSS) for a single process from /proc/{pid}/statm
+     */
+    Status _read_process_memory(pid_t pid, size_t* rss_bytes);
+
+    /**
+     * Refresh memory statistics for all Python processes
+     */
+    void _refresh_memory_stats();
+
     std::unordered_map<PythonVersion, std::vector<ProcessPtr>> _process_pools;
     // Protects _process_pools access
     std::mutex _pools_mutex;
@@ -69,6 +93,7 @@ class PythonServerManager {
     std::atomic<bool> _shutdown_flag {false};
     std::condition_variable _health_check_cv;
     std::mutex _health_check_mutex;
+    MemTracker _mem_tracker {"PythonUDFProcesses"};
 };
 
 } // namespace doris