[Enhancement](pyudf) Enhance Python process poll init and repair

be/src/udf/python/python_server.h

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <atomic>
+#include <chrono>
 #include <condition_variable>
 #include <memory>
 #include <mutex>
@@ -47,7 +48,7 @@ class PythonServerManager {
  std::shared_ptr<T>* client,
  const std::shared_ptr<arrow::Schema>& data_schema = nullptr);
 
- Status fork(const PythonVersion& version, ProcessPtr* process);
+static Status fork(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);
@@ -64,7 +65,11 @@ class PythonServerManager {
  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);
+ std::vector<ProcessPtr> process_pool_snapshot_for_test(const PythonVersion& version);
+
+ bool process_pool_is_initializing_for_test(const PythonVersion& version);
+
+ bool process_pool_is_initialized_for_test(const PythonVersion& version);
 
  Status broadcast_action_to_processes_for_test(const std::string& action_type,
  const std::string& body,
@@ -74,10 +79,28 @@ class PythonServerManager {
 #endif
 
 private:
+ enum class PoolState {
+ UNINITIALIZED,
+ INITIALIZING,
+ // `Initialized` means:
+ // 1. All process slots have attempted initialization.
+ // 2. At least one live process is available for requests.
+ // 3. The health-check thread has been started.
+ INITIALIZED,
+ // Prevent the pool from being incorrectly reused after `shutdown()`
+ STOPPED,
+ };
+
  struct VersionedProcessPool {
  std::mutex mutex;
+ // Coordinates initialization and repair workers with foreground requests.
+ std::condition_variable cv;
  std::vector<ProcessPtr> processes;
- bool initialized = false;
+ PoolState state = PoolState::UNINITIALIZED;
+ // True when at least one process in the pool can serve requests.
+ bool has_available_process = false;
+ // True while a background repair is recreating dead or missing processes.
+ bool repairing = false;
  };
 
  /** 
@@ -104,6 +127,22 @@ class PythonServerManager {
  */
  void _check_and_recreate_processes();
 
+#ifdef BE_TEST
+ static constexpr std::chrono::milliseconds PROCESS_START_TIMEOUT {500};
+ static constexpr std::chrono::milliseconds PROCESS_TERMINATE_TIMEOUT {100};
+ static constexpr std::chrono::milliseconds PROCESS_POOL_INIT_TIMEOUT {1000};
+ static constexpr std::chrono::milliseconds PROCESS_REPAIR_WAIT_TIMEOUT {200};
+#else
+ static constexpr std::chrono::milliseconds PROCESS_START_TIMEOUT {5000};
+ static constexpr std::chrono::milliseconds PROCESS_TERMINATE_TIMEOUT {1000};
+ // FE's default send-fragments RPC timeout is 30s. Keep BE's Python pool wait below it so the
+ // caller sees SERVICE_UNAVAILABLE with Python context instead of a generic RPC deadline error.
+ static constexpr std::chrono::milliseconds PROCESS_POOL_INIT_TIMEOUT {20000};
+ static constexpr std::chrono::milliseconds PROCESS_REPAIR_WAIT_TIMEOUT {1000};
+#endif
+ static int _repair_process_pool(const PythonVersion& version,
+ const std::shared_ptr<VersionedProcessPool>& versioned_pool);
+
  /**
  * Read resident set size (RSS) for a single process from /proc/{pid}/statm
  */
@@ -114,11 +153,14 @@ class PythonServerManager {
  */
  void _refresh_memory_stats();
 
- std::shared_ptr<VersionedProcessPool> _get_or_create_process_pool(const PythonVersion& version);
+ Result<std::shared_ptr<VersionedProcessPool>> _get_or_create_process_pool(
+ const PythonVersion& version);
  std::vector<std::pair<PythonVersion, std::shared_ptr<VersionedProcessPool>>>
  _snapshot_process_pools();
  Status _broadcast_action_to_processes(const std::string& action_type, const std::string& body,
  const std::string& log_name);
+ static bool _select_alive_process_from_pool(const std::vector<ProcessPtr>& pool,
+ ProcessPtr* process);
 
  std::unordered_map<PythonVersion, std::shared_ptr<VersionedProcessPool>> _process_pools;
  // Protects the version -> pool handle map only. Per-version process operations are guarded

be/src/udf/python/python_udf_runtime.cpp

@@ -18,15 +18,177 @@
 #include "udf/python/python_udf_runtime.h"
 
 #include <butil/fd_utility.h>
+#include <signal.h>
+#include <string.h>
 #include <sys/wait.h>
 #include <unistd.h>
 
+#include <algorithm>
 #include <boost/process.hpp>
+#include <cerrno>
+#include <chrono>
+#include <condition_variable>
+#include <deque>
+#include <mutex>
+#include <thread>
+#ifdef BE_TEST
+#include <atomic>
+#endif
 
 #include "common/logging.h"
+#include "runtime/thread_context.h"
 
 namespace doris {
 
+#ifdef BE_TEST
+static constexpr std::chrono::milliseconds PROCESS_TERMINATE_TIMEOUT {100};
+#else
+static constexpr std::chrono::milliseconds PROCESS_TERMINATE_TIMEOUT {1000};
+#endif
+static constexpr std::chrono::milliseconds BACKGROUND_REAP_INTERVAL {1000};
+
+#ifdef BE_TEST
+static std::atomic<int> FORCED_CHILD_EXIT_TIMEOUTS {0};
+
+static bool consume_forced_child_exit_timeout() {
+ int remaining = FORCED_CHILD_EXIT_TIMEOUTS.load(std::memory_order_relaxed);
+ while (remaining > 0) {
+ if (FORCED_CHILD_EXIT_TIMEOUTS.compare_exchange_weak(remaining, remaining - 1,
+ std::memory_order_relaxed)) {
+ return true;
+ }
+ }
+ return false;
+}
+#endif
+
+struct BackgroundChildReaper {
+ std::mutex mutex;
+ std::condition_variable cv;
+ std::deque<pid_t> pids;
+#ifdef BE_TEST
+ std::deque<pid_t> reaped_pids;
+#endif
+ std::thread thread;
+};
+
+static BackgroundChildReaper& background_child_reaper() {
+ static auto* reaper = new BackgroundChildReaper();
+ return *reaper;
+}
+
+void PythonUDFProcess::enqueue_child_for_reap(pid_t pid) {
+ if (pid <= 0) [[unlikely]] {
+ return;
+ }
+
+ auto& reaper = background_child_reaper();
+ {
+ std::lock_guard<std::mutex> lock(reaper.mutex);
+ if (std::find(reaper.pids.begin(), reaper.pids.end(), pid) != reaper.pids.end()) {
+ return;
+ }
+ reaper.pids.push_back(pid);
+ if (!reaper.thread.joinable()) {
+ // This thread only owns pids that were already SIGKILLed but could not be reaped within
+ // the bounded shutdown wait. Such processes may be stuck in uninterruptible I/O; if they
+ // exit later and nobody calls waitpid(), they stay as zombies under BE. Keep reaping them
+ // asynchronously so foreground shutdown remains bounded without dropping wait ownership.
+ reaper.thread = std::thread([]() {
+ SCOPED_INIT_THREAD_CONTEXT();
+ std::deque<pid_t> pending_pids;
+ while (true) {
+ auto& reaper_ref = background_child_reaper();
+ std::unique_lock<std::mutex> lock(reaper_ref.mutex);
+ if (pending_pids.empty()) {
+ reaper_ref.cv.wait(lock,
+ [&reaper_ref]() { return !reaper_ref.pids.empty(); });
+ } else {
+ reaper_ref.cv.wait_for(lock, BACKGROUND_REAP_INTERVAL);
+ }
+ pending_pids.insert(pending_pids.end(), reaper_ref.pids.begin(),
+ reaper_ref.pids.end());
+ reaper_ref.pids.clear();
+ std::deque<pid_t> pids;
+ pids.swap(pending_pids);
+ lock.unlock();
+
+ for (pid_t pending_pid : pids) {
+ int exit_status = 0;
+ auto wait_result = PythonUDFProcess::wait_child_exit(
+ pending_pid, std::chrono::milliseconds(0), &exit_status);
+ if (wait_result == ChildExitWaitResult::EXITED ||
+ wait_result == ChildExitWaitResult::ALREADY_REAPED) {
+ LOG(INFO) << "Background reaped Python process pid=" << pending_pid;
+#ifdef BE_TEST
+ {
+ std::lock_guard<std::mutex> reaped_lock(reaper_ref.mutex);
+ reaper_ref.reaped_pids.push_back(pending_pid);
+ }
+ reaper_ref.cv.notify_all();
+#endif
+ } else if (wait_result == ChildExitWaitResult::TIMEOUT) {
+ pending_pids.push_back(pending_pid);
+ } else {
+ LOG(WARNING) << "Background failed to reap Python process pid="
+ << pending_pid;
+ }
+ }
+ }
+ });
+ }
+ }
+ reaper.cv.notify_one();
+}
+
+#ifdef BE_TEST
+bool PythonUDFProcess::wait_background_reaped_for_test(pid_t pid,
+ std::chrono::milliseconds timeout) {
+ auto& reaper = background_child_reaper();
+ std::unique_lock<std::mutex> lock(reaper.mutex);
+ return reaper.cv.wait_for(lock, timeout, [&reaper, pid]() {
+ return std::find(reaper.reaped_pids.begin(), reaper.reaped_pids.end(), pid) !=
+ reaper.reaped_pids.end();
+ });
+}
+
+void PythonUDFProcess::force_child_exit_timeouts_for_test(int count) {
+ FORCED_CHILD_EXIT_TIMEOUTS.store(count, std::memory_order_relaxed);
+}
+#endif
+
+PythonUDFProcess::ChildExitWaitResult PythonUDFProcess::wait_child_exit(
+ pid_t pid, std::chrono::milliseconds timeout, int* exit_status) {
+#ifdef BE_TEST
+ if (consume_forced_child_exit_timeout()) {
+ return ChildExitWaitResult::TIMEOUT;
+ }
+#endif
+ const auto deadline = std::chrono::steady_clock::now() + timeout;
+ while (true) {
+ pid_t ret = waitpid(pid, exit_status, WNOHANG);
+ if (ret == pid) {
+ return ChildExitWaitResult::EXITED;
+ }
+ if (ret < 0) {
+ if (errno == EINTR) {
+ // retry if interrupted
+ continue;
+ }
+ // Another owner may already have observed the child exit through boost::process.
+ if (errno == ECHILD) {
+ return ChildExitWaitResult::ALREADY_REAPED;
+ }
+ LOG(WARNING) << "Failed to wait Python process pid=" << pid << ": " << strerror(errno);
+ return ChildExitWaitResult::ERROR;
+ }
+ if (std::chrono::steady_clock::now() >= deadline) {
+ return ChildExitWaitResult::TIMEOUT;
+ }
+ std::this_thread::sleep_for(std::chrono::milliseconds(10));
+ }
+}
+
 void PythonUDFProcess::remove_unix_socket() {
  if (_uri.empty() || _unix_socket_file_path.empty()) return;
 
@@ -47,29 +209,51 @@ void PythonUDFProcess::remove_unix_socket() {
 void PythonUDFProcess::shutdown() {
  if (!_child.valid() || _is_shutdown) return;
 
- _child.terminate();
- bool graceful = false;
- constexpr std::chrono::milliseconds retry_interval(100); // 100ms
-
- for (int i = 0; i < TERMINATE_RETRY_TIMES; ++i) {
- if (!_child.running()) {
- graceful = true;
- break;
- }
- std::this_thread::sleep_for(retry_interval);
+ int exit_status = 0;
+ bool exited = !_child.running();
+ bool status_available = false;
+ bool already_reaped = false;
+ if (!exited) {
+ ::kill(_child_pid, SIGTERM);
+ auto wait_result = wait_child_exit(_child_pid, PROCESS_TERMINATE_TIMEOUT, &exit_status);
+ exited = wait_result == ChildExitWaitResult::EXITED ||
+ wait_result == ChildExitWaitResult::ALREADY_REAPED;
+ status_available = wait_result == ChildExitWaitResult::EXITED;
+ already_reaped = wait_result == ChildExitWaitResult::ALREADY_REAPED;
+ } else {
+ auto wait_result = wait_child_exit(_child_pid, std::chrono::milliseconds(0), &exit_status);
+ status_available = wait_result == ChildExitWaitResult::EXITED;
+ already_reaped = wait_result == ChildExitWaitResult::ALREADY_REAPED;
  }
 
- if (!graceful) {
- LOG(WARNING) << "Python process did not terminate gracefully, sending SIGKILL";
+ if (!exited) {
+ LOG(WARNING) << "Python process did not terminate gracefully, sending SIGKILL, pid="
+ << _child_pid;
  ::kill(_child_pid, SIGKILL);
- _child.wait();
+ auto wait_result = wait_child_exit(_child_pid, PROCESS_TERMINATE_TIMEOUT, &exit_status);
+ exited = wait_result == ChildExitWaitResult::EXITED ||
+ wait_result == ChildExitWaitResult::ALREADY_REAPED;
+ status_available = wait_result == ChildExitWaitResult::EXITED;
+ already_reaped = wait_result == ChildExitWaitResult::ALREADY_REAPED;
  }
+ _child.detach();
 
- if (int exit_code = _child.exit_code(); exit_code > 128 && exit_code <= 255) {
- int signal = exit_code - 128;
- LOG(INFO) << "Python process was killed by signal " << signal;
+ if (!exited) [[unlikely]] {
+ LOG(WARNING) << "Python process did not exit after SIGKILL, enqueue background reap, pid="
+ << _child_pid;
+ enqueue_child_for_reap(_child_pid);
+ } else if (already_reaped) {
+ LOG(INFO) << "Python process already reaped by another owner, pid=" << _child_pid;
+ } else if (!status_available) {
+ LOG(INFO) << "Python process exited but exit status is unavailable, pid=" << _child_pid;
  } else {
- LOG(INFO) << "Python process exited normally with code: " << exit_code;
+ if (WIFSIGNALED(exit_status)) {
+ LOG(INFO) << "Python process was killed by signal " << WTERMSIG(exit_status);
+ } else if (WIFEXITED(exit_status)) {
+ LOG(INFO) << "Python process exited normally with code: " << WEXITSTATUS(exit_status);
+ } else {
+ LOG(INFO) << "Python process exited";
+ }
  }
 
  _output_stream.close();
@@ -84,4 +268,4 @@ std::string PythonUDFProcess::to_string() const {
  _child_pid, _uri, _unix_socket_file_path, _is_shutdown);
 }
 
-} // namespace doris
+} // namespace doris

be/src/udf/python/python_udf_runtime.h

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <boost/process.hpp>
+#include <chrono>
 
 #include "udf/python/python_env.h"
 
@@ -75,6 +76,15 @@ class PythonUDFProcess {
 
  void shutdown();
 
+ enum class ChildExitWaitResult { EXITED, ALREADY_REAPED, TIMEOUT, ERROR };
+
+ static ChildExitWaitResult wait_child_exit(pid_t pid, std::chrono::milliseconds timeout,
+ int* exit_status);
+
+ // Hand off a killed child that could not be reaped synchronously. The background reaper keeps
+ // waitpid ownership so a later child exit will not become a zombie under BE.
+ static void enqueue_child_for_reap(pid_t pid);
+
  std::string to_string() const;
 
  pid_t get_child_pid() const { return _child_pid; }
@@ -85,10 +95,13 @@ class PythonUDFProcess {
 
 #ifdef BE_TEST
  void set_uri_for_test(std::string uri) { _uri = std::move(uri); }
+
+ static bool wait_background_reaped_for_test(pid_t pid, std::chrono::milliseconds timeout);
+
+ static void force_child_exit_timeouts_for_test(int count);
 #endif
 
 private:
- constexpr static int TERMINATE_RETRY_TIMES = 10;
  constexpr static size_t MAX_ACCUMULATED_LOG_SIZE = 65536;
 
  bool _is_shutdown {false};