[Fix](pyudf) make Python server pool selection alive-aware and version-isolated (#62620)

`PythonServerManager` does not check if the python process corresponding
to the version is alive when retrieving the process, which may cause
errors like:
```text
java.lang.IllegalStateException: PYTHON_UDF_BLOCKED suite=python_udf_cross_feature_import_storage scenario=python_udf_cross_feature_import_storage.inline_runtime reason=inline probe failed. reason=errCode = 2, detailMessage = (172.20.49.73)[INTERNAL_ERROR]IOError: Flight stream finish failed with gRPC code 14, message: failed to connect to all addresses; last error: UNKNOWN: unix:/tmp/doris_python_udf_55799.sock: Connection refused
```

After modification, before obtaining the Python process, it will check
if the process is alive to ensure the availability of this feature.

Author: linrrzqqq

be/src/udf/python/python_server.cpp

@@ -27,6 +27,7 @@
 #include <boost/asio.hpp>
 #include <boost/process.hpp>
 #include <fstream>
+#include <future>
 
 #include "arrow/flight/client.h"
 #include "common/config.h"
@@ -37,32 +38,70 @@
 
 namespace doris {
 
-template <typename T>
+std::shared_ptr<PythonServerManager::VersionedProcessPool>
+PythonServerManager::_get_or_create_process_pool(const PythonVersion& version) {
+    std::lock_guard<std::mutex> lock(_pools_mutex);
+    auto& pool = _process_pools[version];
+    if (!pool) {
+        pool = std::make_shared<VersionedProcessPool>();
+    }
+    return pool;
+}
+
+std::vector<std::pair<PythonVersion, std::shared_ptr<PythonServerManager::VersionedProcessPool>>>
+PythonServerManager::_snapshot_process_pools() {
+    std::lock_guard<std::mutex> lock(_pools_mutex);
+    std::vector<std::pair<PythonVersion, std::shared_ptr<VersionedProcessPool>>> snapshot;
+    snapshot.reserve(_process_pools.size());
+    for (const auto& [version, pool] : _process_pools) {
+        snapshot.emplace_back(version, pool);
+    }
+    return snapshot;
+}
+
+#ifdef BE_TEST
+void PythonServerManager::set_process_pool_for_test(const PythonVersion& version,
+                                                    std::vector<ProcessPtr> processes,
+                                                    bool initialized) {
+    auto versioned_pool = _get_or_create_process_pool(version);
+    std::lock_guard<std::mutex> lock(versioned_pool->mutex);
+    versioned_pool->processes = std::move(processes);
+    versioned_pool->initialized = initialized;
+}
+
+std::vector<ProcessPtr>& PythonServerManager::process_pool_for_test(const PythonVersion& version) {
+    auto versioned_pool = _get_or_create_process_pool(version);
+    return versioned_pool->processes;
+}
+#endif
+
+template <typename ClientType>
 Status PythonServerManager::get_client(const PythonUDFMeta& func_meta, const PythonVersion& version,
-                                       std::shared_ptr<T>* client,
+                                       std::shared_ptr<ClientType>* client,
                                        const std::shared_ptr<arrow::Schema>& data_schema) {
-    // Ensure process pool is initialized for this version
-    RETURN_IF_ERROR(ensure_pool_initialized(version));
+    std::shared_ptr<VersionedProcessPool> versioned_pool =
+            DORIS_TRY(_ensure_pool_initialized(version));
 
     ProcessPtr process;
-    RETURN_IF_ERROR(get_process(version, &process));
+    RETURN_IF_ERROR(_get_process(version, versioned_pool, &process));
 
-    if constexpr (std::is_same_v<T, PythonUDAFClient>) {
-        RETURN_IF_ERROR(T::create(func_meta, std::move(process), data_schema, client));
+    if constexpr (std::is_same_v<ClientType, PythonUDAFClient>) {
+        RETURN_IF_ERROR(ClientType::create(func_meta, std::move(process), data_schema, client));
     } else {
-        RETURN_IF_ERROR(T::create(func_meta, std::move(process), client));
+        RETURN_IF_ERROR(ClientType::create(func_meta, std::move(process), client));
     }
 
     return Status::OK();
 }
 
-Status PythonServerManager::ensure_pool_initialized(const PythonVersion& version) {
-    std::lock_guard<std::mutex> lock(_pools_mutex);
+Result<std::shared_ptr<PythonServerManager::VersionedProcessPool>>
+PythonServerManager::_ensure_pool_initialized(const PythonVersion& version) {
+    auto versioned_pool = _get_or_create_process_pool(version);
+    std::lock_guard<std::mutex> lock(versioned_pool->mutex);
 
     // Check if already initialized
-    if (_initialized_versions.count(version)) return Status::OK();
+    if (versioned_pool->initialized) return versioned_pool;
 
-    std::vector<ProcessPtr>& pool = _process_pools[version];
     // 0 means use CPU core count as default, otherwise use the specified value
     int max_pool_size = config::max_python_process_num > 0 ? config::max_python_process_num
                                                            : CpuInfo::num_cores();
@@ -91,7 +130,7 @@ Status PythonServerManager::ensure_pool_initialized(const PythonVersion& version
     for (int i = 0; i < max_pool_size; i++) {
         Status s = futures[i].get();
         if (s.ok() && temp_processes[i]) {
-            pool.push_back(std::move(temp_processes[i]));
+            versioned_pool->processes.emplace_back(std::move(temp_processes[i]));
             success_count++;
         } else {
             failure_count++;
@@ -100,38 +139,64 @@ Status PythonServerManager::ensure_pool_initialized(const PythonVersion& version
         }
     }
 
-    if (pool.empty()) {
-        return Status::InternalError(
+    if (versioned_pool->processes.empty()) {
+        return ResultError(Status::Error<ErrorCode::SERVICE_UNAVAILABLE>(
                 "Failed to initialize Python process pool: all {} process creation attempts failed",
-                max_pool_size);
+                max_pool_size));
     }
 
     LOG(INFO) << "Python process pool initialized for version " << version.to_string()
               << ": created " << success_count << " processes"
               << (failure_count > 0 ? fmt::format(" ({} failed)", failure_count) : "");
 
-    _initialized_versions.insert(version);
+    versioned_pool->initialized = true;
     _start_health_check_thread();
 
-    return Status::OK();
+    return versioned_pool;
 }
 
-Status PythonServerManager::get_process(const PythonVersion& version, ProcessPtr* process) {
-    std::lock_guard<std::mutex> lock(_pools_mutex);
-    std::vector<ProcessPtr>& pool = _process_pools[version];
+Status PythonServerManager::_get_process(
+        const PythonVersion& version, const std::shared_ptr<VersionedProcessPool>& versioned_pool,
+        ProcessPtr* process) {
+    std::lock_guard<std::mutex> lock(versioned_pool->mutex);
+    std::vector<ProcessPtr>& pool = versioned_pool->processes;
 
     if (UNLIKELY(pool.empty())) {
         return Status::InternalError("Python process pool is empty for version {}",
                                      version.to_string());
     }
 
-    // Find process with minimum load (use_count - 1 gives active client count)
-    auto min_iter = std::min_element(
-            pool.begin(), pool.end(),
-            [](const ProcessPtr& a, const ProcessPtr& b) { return a.use_count() < b.use_count(); });
+    // Prefer an already-alive process and only use load balancing inside that alive subset.
+    // keep dead entries stay in the pool for the background health checker
+    // unless there is no alive process left for the current request.
+    auto min_alive_iter = std::min_element(pool.begin(), pool.end(),
+                                           [](const ProcessPtr& a, const ProcessPtr& b) {
+                                               const bool a_alive = a && a->is_alive();
+                                               const bool b_alive = b && b->is_alive();
+                                               if (a_alive != b_alive) {
+                                                   return a_alive > b_alive;
+                                               }
+                                               return a.use_count() < b.use_count();
+                                           });
+
+    if (min_alive_iter != pool.end() && *min_alive_iter && (*min_alive_iter)->is_alive()) {
+        *process = *min_alive_iter;
+        return Status::OK();
+    }
+
+    // Only reach here when the pool has no alive process at all. Try one foreground
+    // recovery so the caller has a chance to proceed; leave batch repair to health check.
+    auto& candidate = pool.front();
+    ProcessPtr replacement;
+    Status status = fork(version, &replacement);
+    if (!status.ok()) {
+        return Status::Error<ErrorCode::SERVICE_UNAVAILABLE>(
+                "Python process pool has no available process for version {}, reason: {}",
+                version.to_string(), status.to_string());
+    }
 
-    // Return process with minimum load
-    *process = *min_iter;
+    candidate = std::move(replacement);
+    *process = candidate;
     return Status::OK();
 }
 
@@ -191,6 +256,7 @@ Status PythonServerManager::fork(const PythonVersion& version, ProcessPtr* proce
 }
 
 void PythonServerManager::_start_health_check_thread() {
+    std::lock_guard<std::mutex> lock(_health_check_mutex);
     if (_health_check_thread) return;
 
     LOG(INFO) << "Starting Python process health check thread (interval: 30 seconds)";
@@ -217,13 +283,13 @@ void PythonServerManager::_start_health_check_thread() {
 }
 
 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 (auto& [version, versioned_pool] : _snapshot_process_pools()) {
+        std::lock_guard<std::mutex> lock(versioned_pool->mutex);
+        auto& pool = versioned_pool->processes;
         for (size_t i = 0; i < pool.size(); ++i) {
             auto& process = pool[i];
             if (!process) continue;
@@ -268,15 +334,22 @@ void PythonServerManager::shutdown() {
     }
 
     // Shutdown all processes
-    std::lock_guard<std::mutex> lock(_pools_mutex);
-    for (auto& [version, pool] : _process_pools) {
+    for (auto& [version, versioned_pool] : _snapshot_process_pools()) {
+        std::lock_guard<std::mutex> lock(versioned_pool->mutex);
+        auto& pool = versioned_pool->processes;
         for (auto& process : pool) {
             if (process) {
                 process->shutdown();
             }
         }
+        pool.clear();
+        versioned_pool->initialized = false;
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(_pools_mutex);
+        _process_pools.clear();
     }
-    _process_pools.clear();
 }
 
 Status PythonServerManager::_read_process_memory(pid_t pid, size_t* rss_bytes) {
@@ -305,11 +378,11 @@ Status PythonServerManager::_read_process_memory(pid_t pid, size_t* rss_bytes) {
 }
 
 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& [version, versioned_pool] : _snapshot_process_pools()) {
+        std::lock_guard<std::mutex> lock(versioned_pool->mutex);
+        const auto& pool = versioned_pool->processes;
         for (const auto& process : pool) {
             if (!process || !process->is_alive()) continue;
 
@@ -339,15 +412,15 @@ Status PythonServerManager::clear_module_cache(const std::string& location) {
         return Status::InvalidArgument("Empty location for clear_module_cache");
     }
 
-    std::lock_guard<std::mutex> lock(_pools_mutex);
-
     std::string body = fmt::format(R"({{"location": "{}"}})", location);
 
     int success_count = 0;
     int fail_count = 0;
     bool has_active_process = false;
 
-    for (auto& [version, pool] : _process_pools) {
+    for (auto& [version, versioned_pool] : _snapshot_process_pools()) {
+        std::lock_guard<std::mutex> lock(versioned_pool->mutex);
+        auto& pool = versioned_pool->processes;
         for (auto& process : pool) {
             if (!process || !process->is_alive()) {
                 continue;
@@ -422,4 +495,4 @@ template Status PythonServerManager::get_client<PythonUDTFClient>(
         std::shared_ptr<PythonUDTFClient>* client,
         const std::shared_ptr<arrow::Schema>& data_schema);
 
-} // namespace doris
+} // namespace doris

be/src/udf/python/python_server.h

@@ -20,7 +20,10 @@
 #include <atomic>
 #include <condition_variable>
 #include <memory>
+#include <mutex>
 #include <thread>
+#include <unordered_map>
+#include <vector>
 
 #include "common/status.h"
 #include "runtime/memory/mem_tracker.h"
@@ -46,25 +49,41 @@ class PythonServerManager {
 
     Status fork(const PythonVersion& version, ProcessPtr* process);
 
-    Status get_process(const PythonVersion& version, ProcessPtr* process);
-
     // Clear Python module cache for a specific UDF location across all processes
     Status clear_module_cache(const std::string& location);
 
-    Status ensure_pool_initialized(const PythonVersion& version);
-
     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;
-    }
+    void set_process_pool_for_test(const PythonVersion& version, std::vector<ProcessPtr> processes,
+                                   bool initialized = true);
+
+    std::vector<ProcessPtr>& process_pool_for_test(const PythonVersion& version);
 #endif
 
 private:
+    struct VersionedProcessPool {
+        std::mutex mutex;
+        std::vector<ProcessPtr> processes;
+        bool initialized = false;
+    };
+
+    /** 
+     * Lazily initialize and return the process pool for specific Python version. 
+     */
+    Result<std::shared_ptr<VersionedProcessPool>> _ensure_pool_initialized(
+            const PythonVersion& version);
+
+    /**
+     * Pick an available process from specific pool, recreating one on demand if needed.
+     */
+    Status _get_process(const PythonVersion& version,
+                        const std::shared_ptr<VersionedProcessPool>& versioned_pool,
+                        ProcessPtr* process);
+
     /**
      * Start health check background thread (called once by ensure_pool_initialized)
      * Thread periodically checks process health and refreshes memory stats
@@ -86,11 +105,14 @@ class PythonServerManager {
      */
     void _refresh_memory_stats();
 
-    std::unordered_map<PythonVersion, std::vector<ProcessPtr>> _process_pools;
-    // Protects _process_pools access
+    std::shared_ptr<VersionedProcessPool> _get_or_create_process_pool(const PythonVersion& version);
+    std::vector<std::pair<PythonVersion, std::shared_ptr<VersionedProcessPool>>>
+    _snapshot_process_pools();
+
+    std::unordered_map<PythonVersion, std::shared_ptr<VersionedProcessPool>> _process_pools;
+    // Protects the version -> pool handle map only. Per-version process operations are guarded
+    // by VersionedProcessPool::mutex.
     std::mutex _pools_mutex;
-    // Track which versions have been initialized
-    std::unordered_set<PythonVersion> _initialized_versions;
     // Health check background thread
     std::unique_ptr<std::thread> _health_check_thread;
     std::atomic<bool> _shutdown_flag {false};
@@ -99,4 +121,4 @@ class PythonServerManager {
     MemTracker _mem_tracker {"PythonUDFProcesses"};
 };
 
-} // namespace doris
+} // namespace doris