[refactor](udf) Isolate Python process pools by function id

be/src/agent/task_worker_pool.cpp

@@ -59,6 +59,7 @@
 #include "common/metrics/doris_metrics.h"
 #include "common/status.h"
 #include "cpp/obj-client/obj_storage_client.h"
+#include "cpp/sync_point.h"
 #include "exec/spill/spill_file_manager.h"
 #include "io/fs/file_system.h"
 #include "io/fs/hdfs_file_system.h"
@@ -154,6 +155,7 @@ void remove_task_info(const TTaskType::type task_type, int64_t signature) {
 }
 
 void finish_task(const TFinishTaskRequest& finish_task_request) {
+    TEST_SYNC_POINT_RETURN_WITH_VOID("TaskWorkerPool::finish_task", &finish_task_request);
     // Return result to FE
     TMasterResult result;
     uint32_t try_time = 0;
@@ -2579,34 +2581,50 @@ void clean_trash_callback(StorageEngine& engine, const TAgentTaskRequest& req) {
 }
 
 void clean_udf_cache_callback(const TAgentTaskRequest& req) {
+    Status status;
     const auto& clean_req = req.clean_udf_cache_req;
     if (clean_req.__isset.function_id && clean_req.function_id <= 0) {
         LOG(WARNING) << "skip clean udf cache request with invalid function_id="
                      << clean_req.function_id
                      << ", function_signature=" << clean_req.function_signature;
-        return;
-    }
-    // Requests from old FEs do not set function_id and must keep signature-based cleanup.
-    const bool drop_by_function_id = clean_req.__isset.function_id;
-
-    if (doris::config::enable_java_support) {
-        // The id, not just the signature: it is what the java-udf plugin keys its compiled
-        // classes by, because the signature carries no database and FE renders a variadic one
-        // differently here than on the requests that execute the function.
-        WARN_IF_ERROR(
-                Jni::PluginRegistry::clean_udf_cache(
-                        drop_by_function_id ? clean_req.function_id : 0,
-                        clean_req.function_signature),
-                fmt::format("failed to clean Java UDF cache, function_signature={}, function_id={}",
+        status = Status::InvalidArgument("Invalid function_id for UDF cache cleanup: {}",
+                                         clean_req.function_id);
+    } else {
+        // Requests from old FEs do not set function_id and must keep signature-based cleanup.
+        const bool drop_by_function_id = clean_req.__isset.function_id;
+
+        if (doris::config::enable_java_support) {
+            // The id, not just the signature: it is what the java-udf plugin keys its compiled
+            // classes by, because the signature carries no database and FE renders a variadic one
+            // differently here than on the requests that execute the function.
+            WARN_IF_ERROR(
+                    Jni::PluginRegistry::clean_udf_cache(
+                            drop_by_function_id ? clean_req.function_id : 0,
+                            clean_req.function_signature),
+                    fmt::format(
+                            "failed to clean Java UDF cache, function_signature={}, function_id={}",
                             clean_req.function_signature, clean_req.function_id));
-    }
-    if (drop_by_function_id) {
-        UserFunctionCache::instance()->drop_function_cache(clean_req.function_id);
-        PythonServerManager::instance().clear_udaf_state_cache(clean_req.function_id);
+        }
+        if (drop_by_function_id) {
+            UserFunctionCache::instance()->drop_function_cache(clean_req.function_id);
+            PythonServerManager::instance().drop_function(clean_req.function_id);
+        }
+
+        LOG(INFO) << "clean udf cache callback finish: function_signature="
+                  << clean_req.function_signature << ", function_id=" << clean_req.function_id;
     }
 
-    LOG(INFO) << "clean udf cache callback finish: function_signature="
-              << clean_req.function_signature << ", function_id=" << clean_req.function_id;
+    // Report the result only for queued tasks; one-shot tasks use signature -1.
+    if (req.signature != -1) {
+        TFinishTaskRequest finish_task_request;
+        finish_task_request.__set_backend(BackendOptions::get_local_backend());
+        finish_task_request.__set_task_type(req.task_type);
+        finish_task_request.__set_signature(req.signature);
+        finish_task_request.__set_task_status(status.to_thrift());
+        finish_task(finish_task_request);
+        // Stop reporting this task as running so a failed result submission can trigger a retry.
+        remove_task_info(req.task_type, req.signature);
+    }
 }
 
 void report_index_policy_callback(const ClusterInfo* cluster_info) {

be/src/common/config.cpp

@@ -1198,9 +1198,6 @@ DEFINE_String(python_conda_root_path, "");
 DEFINE_String(python_venv_root_path, "${DORIS_HOME}/lib/udf/python");
 // python interpreter paths used by venv, e.g. /usr/bin/python3.7:/usr/bin/python3.6
 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");

be/src/common/config.h

@@ -1271,8 +1271,6 @@ DECLARE_String(python_conda_root_path);
 DECLARE_String(python_venv_root_path);
 // python interpreter paths used by venv, e.g. /usr/bin/python3.7:/usr/bin/python3.6
 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);
 

be/src/exprs/aggregate/aggregate_function_python_udaf.cpp

@@ -191,6 +191,10 @@ Status AggregatePythonUDAFData::get(IColumn& to, const DataTypePtr& result_type,
 Status AggregatePythonUDAF::open() {
     // Build function metadata from TFunction
     _func_meta.id = _fn.id;
+    _func_meta.pool_size = _fn.__isset.pool_size ? _fn.pool_size : 3;
+    _func_meta.auto_suspend = _fn.__isset.auto_suspend
+                                      ? _fn.auto_suspend
+                                      : PythonUDFMeta::DEFAULT_AUTO_SUSPEND_SECONDS;
     _func_meta.name = _fn.name.function_name;
 
     // For UDAF, symbol is in aggregate_fn
@@ -236,7 +240,8 @@ Status AggregatePythonUDAF::open() {
 
     if (_func_meta.type == PythonUDFLoadType::MODULE) {
         RETURN_IF_ERROR(UserFunctionCache::instance()->get_pypath(
-                _func_meta.id, _func_meta.location, _func_meta.checksum, &_func_meta.location));
+                _func_meta.id, _func_meta.location, _func_meta.checksum, &_func_meta.location,
+                &_package_cache_entry));
     }
 
     return Status::OK();
@@ -272,7 +277,7 @@ void AggregatePythonUDAF::create(AggregateDataPtr __restrict place) const {
     DCHECK(reinterpret_cast<Data*>(place)) << "Place must not be null";
 
     if (Status st = PythonServerManager::instance().get_client(
-                _func_meta, _python_version, &(this->data(place).client), _schema);
+                _func_meta, _python_version, &data(place).client, _schema, _package_cache_entry);
         UNLIKELY(!st.ok())) {
         this->data(place).~Data();
         throw doris::Exception(ErrorCode::INTERNAL_ERROR, "Failed to get Python UDAF client: {}",

be/src/exprs/aggregate/aggregate_function_python_udaf.h

@@ -28,6 +28,8 @@
 
 namespace doris {
 
+struct UserFunctionCacheEntry;
+
 /**
  * Aggregate state data for Python UDAF
  * 
@@ -172,6 +174,7 @@ class AggregatePythonUDAF final
     DataTypePtr _return_type;
     PythonUDFMeta _func_meta;
     PythonVersion _python_version;
+    std::shared_ptr<UserFunctionCacheEntry> _package_cache_entry;
     // Arrow Flight schema: [argument_types..., places: int64, binary_data: binary]
     // Used for all UDAF RPC operations
     // - places column is always present (NULL in single-place mode, actual place_id values in GROUP BY mode)

be/src/exprs/function/function_python_udf.cpp

@@ -53,6 +53,9 @@ Status PythonFunctionCall::open(FunctionContext* context,
     PythonVersion version;
     PythonUDFMeta func_meta;
     func_meta.id = _fn.id;
+    func_meta.pool_size = _fn.__isset.pool_size ? _fn.pool_size : 3;
+    func_meta.auto_suspend = _fn.__isset.auto_suspend ? _fn.auto_suspend
+                                                      : PythonUDFMeta::DEFAULT_AUTO_SUSPEND_SECONDS;
     func_meta.name = _fn.name.function_name;
     func_meta.symbol = _fn.scalar_fn.symbol;
     if (!_fn.function_code.empty()) {
@@ -85,13 +88,16 @@ Status PythonFunctionCall::open(FunctionContext* context,
     LOG(INFO) << fmt::format("runtime_version: {}, func_meta: {}", version.to_string(),
                              func_meta.to_string());
 
+    std::shared_ptr<UserFunctionCacheEntry> package_cache_entry;
     if (func_meta.type == PythonUDFLoadType::MODULE) {
         RETURN_IF_ERROR(UserFunctionCache::instance()->get_pypath(
-                func_meta.id, func_meta.location, func_meta.checksum, &func_meta.location));
+                func_meta.id, func_meta.location, func_meta.checksum, &func_meta.location,
+                &package_cache_entry));
     }
 
     PythonUDFClientPtr client = nullptr;
-    RETURN_IF_ERROR(PythonServerManager::instance().get_client(func_meta, version, &client));
+    RETURN_IF_ERROR(PythonServerManager::instance().get_client(func_meta, version, &client, nullptr,
+                                                               std::move(package_cache_entry)));
 
     if (!client) {
         return Status::InternalError("Python UDF client is null");

be/src/exprs/table_function/python_udtf_function.cpp

@@ -61,6 +61,10 @@ PythonUDTFFunction::PythonUDTFFunction(const TFunction& t_fn) : TableFunction(),
 Status PythonUDTFFunction::open() {
     PythonUDFMeta python_udf_meta;
     python_udf_meta.id = _t_fn.id;
+    python_udf_meta.pool_size = _t_fn.__isset.pool_size ? _t_fn.pool_size : 3;
+    python_udf_meta.auto_suspend = _t_fn.__isset.auto_suspend
+                                           ? _t_fn.auto_suspend
+                                           : PythonUDFMeta::DEFAULT_AUTO_SUSPEND_SECONDS;
     python_udf_meta.name = _t_fn.name.function_name;
     python_udf_meta.symbol = _t_fn.scalar_fn.symbol;
 
@@ -79,10 +83,11 @@ Status PythonUDTFFunction::open() {
 
     python_udf_meta.client_type = PythonClientType::UDTF;
 
+    std::shared_ptr<UserFunctionCacheEntry> package_cache_entry;
     if (python_udf_meta.type == PythonUDFLoadType::MODULE) {
         RETURN_IF_ERROR(UserFunctionCache::instance()->get_pypath(
                 python_udf_meta.id, python_udf_meta.location, python_udf_meta.checksum,
-                &python_udf_meta.location));
+                &python_udf_meta.location, &package_cache_entry));
     }
 
     PythonVersion version;
@@ -105,8 +110,8 @@ Status PythonUDTFFunction::open() {
     python_udf_meta.always_nullable = python_udf_meta.return_type->is_nullable();
     RETURN_IF_ERROR(python_udf_meta.check());
 
-    RETURN_IF_ERROR(
-            PythonServerManager::instance().get_client(python_udf_meta, version, &_udtf_client));
+    RETURN_IF_ERROR(PythonServerManager::instance().get_client(
+            python_udf_meta, version, &_udtf_client, nullptr, std::move(package_cache_entry)));
 
     if (!_udtf_client) {
         return Status::InternalError("Failed to create Python UDTF client");

be/src/runtime/user_function_cache.cpp

@@ -498,10 +498,14 @@ Status UserFunctionCache::get_jarpath(int64_t fid, const std::string& url,
 }
 
 Status UserFunctionCache::get_pypath(int64_t fid, const std::string& url,
-                                     const std::string& checksum, std::string* libpath) {
+                                     const std::string& checksum, std::string* libpath,
+                                     std::shared_ptr<UserFunctionCacheEntry>* cache_entry) {
     std::shared_ptr<UserFunctionCacheEntry> entry = nullptr;
     RETURN_IF_ERROR(_get_cache_entry(fid, url, checksum, entry, LibType::PY_ZIP));
     *libpath = entry->lib_file;
+    if (cache_entry != nullptr) {
+        *cache_entry = std::move(entry);
+    }
     return Status::OK();
 }
 
@@ -579,17 +583,11 @@ void UserFunctionCache::drop_function_cache(int64_t fid) {
         _entry_map.erase(it);
     }
 
-    // For Python UDF, clear module cache in Python server before deleting files
-    if (entry->type == LibType::PY_ZIP && !entry->lib_file.empty()) {
-        auto status = PythonServerManager::instance().clear_module_cache(entry->lib_file);
-        if (!status.ok()) [[unlikely]] {
-            LOG(WARNING) << "drop_function_cache: failed to clear Python module cache for "
-                         << entry->lib_file << ": " << status.to_string();
-        }
-    }
-
     // Mark for deletion, destructor will delete the files
     entry->should_delete_library.store(true);
+    if (entry->type == LibType::PY_ZIP) {
+        PythonServerManager::instance().drop_function(fid);
+    }
 }
 
 } // namespace doris

be/src/runtime/user_function_cache.h

@@ -59,8 +59,11 @@ class UserFunctionCache {
     Status get_jarpath(int64_t fid, const std::string& url, const std::string& checksum,
                        std::string* libpath);
 
+    // Optionally return a shared cache entry to pin package files: DROP only removes them
+    // after the last query/process releases this handle. The handle is for lifetime only.
     Status get_pypath(int64_t fid, const std::string& url, const std::string& checksum,
-                      std::string* libpath);
+                      std::string* libpath,
+                      std::shared_ptr<UserFunctionCacheEntry>* cache_entry = nullptr);
 
     // Drop the cached function library by function id.
     void drop_function_cache(int64_t fid);

be/src/udf/python/python_client.cpp

@@ -72,22 +72,20 @@ Status PythonClient::init(const PythonUDFMeta& func_meta, ProcessPtr process) {
 }
 
 Status PythonClient::close() {
-    if (!_inited || !_writer) {
-        return Status::OK();
-    }
-
-    auto writer_res = _writer->Close();
-    if (!writer_res.ok()) {
-        // Don't propagate error from close, just log it
-        LOG(WARNING) << "Error closing Python client writer: " << writer_res.message();
+    if (_inited && _writer) {
+        auto writer_res = _writer->Close();
+        if (!writer_res.ok()) {
+            LOG(WARNING) << "Error closing Python client writer: " << writer_res.message();
+        }
     }
-
+    // Error handling may already have reset the writer; still release the pool lease.
     _inited = false;
     _begin = false;
     _arrow_client.reset();
     _writer.reset();
     _reader.reset();
     _process.reset();
+    _pool_lease.reset();
 
     return Status::OK();
 }

be/src/udf/python/python_client.h

@@ -28,6 +28,8 @@
 
 namespace doris {
 
+class PythonUDFPoolLease;
+
 /**
  * Base class for Python UDF/UDAF/UDTF clients
  * 
@@ -102,6 +104,9 @@ class PythonClient {
      */
     Status read_batch(std::shared_ptr<arrow::RecordBatch>* output);
 
+    // Released after the stream and process, so idle reclamation cannot interrupt cleanup.
+    std::shared_ptr<PythonUDFPoolLease> _pool_lease;
+
     // Common state
     bool _inited = false;
     bool _begin = false;         // Track if Begin() has been called
@@ -112,6 +117,8 @@ class PythonClient {
     ProcessPtr _process;
 
 private:
+    friend class PythonServerManager;
+
     DISALLOW_COPY_AND_ASSIGN(PythonClient);
 };