Skip to content

[Fix](python udf) isolate module caches across functions #67511

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
linrrzqqq wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

[Fix](python udf) isolate module caches across functions #67511

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7a984c9
[fix](python udf) Fix module cache cleanup
linrrzqqq Aug 31, 2026
d9b7489
[fix](python udf) Validate cache action results
linrrzqqq Aug 31, 2026
d8b55d1
[fix](python udf) isolate module caches across functions
linrrzqqq Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion be/src/agent/task_worker_pool.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2599,8 +2599,13 @@ void clean_udf_cache_callback(const TAgentTaskRequest& req) {
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);
auto status = UserFunctionCache::instance()->drop_function_cache(clean_req.function_id);
if (!status.ok()) {

@github-actions github-actions Bot Sep 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

drop_function_cache intentionally keeps the entry/files when clear_module_cache fails, but this callback only logs the error and returns from a void worker task. DropFunctionCommand submits CleanUDFCacheTask once, and the task pool has no retry/requeue or completion failure propagation, so a transiently dead Python process permanently leaks the cache entry and both files. Please add a deferred/bounded retry (or another recovery queue) and surface an unrecoverable failure instead of treating this task as handled.

LOG(WARNING) << "failed to drop function cache for function_id="
<< clean_req.function_id << ": " << status.to_string();
}
PythonServerManager::instance().clear_udaf_state_cache(clean_req.function_id);
}

Expand Down
23 changes: 15 additions & 8 deletions be/src/runtime/user_function_cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -567,29 +567,36 @@ Status UserFunctionCache::_check_and_return_default_java_udf_url(const std::stri
return Status::OK();
}

void UserFunctionCache::drop_function_cache(int64_t fid) {
Status UserFunctionCache::drop_function_cache(int64_t fid) {
std::shared_ptr<UserFunctionCacheEntry> entry = nullptr;
{
std::lock_guard<std::mutex> l(_cache_lock);
auto it = _entry_map.find(fid);
if (it == _entry_map.end()) {
return;
return Status::OK();
}
entry = it->second;
_entry_map.erase(it);
}

// lib_file changes from the downloaded zip path to the extracted directory
// while an entry is loaded. Wait for that transition before clearing Python.
std::unique_lock<std::mutex> load_lock(entry->load_lock);

@github-actions github-actions Bot Sep 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

drop_function_cache now leaves the entry in _entry_map while waiting on load_lock and broadcasting the Python clear. A concurrent _get_cache_entry can capture that entry, return its already-loaded lib_file, and then race with the erase/should_delete_library path; once shared references drain, UserFunctionCacheEntry::~UserFunctionCacheEntry deletes the directory/zip even though the caller has just received the path. This can make a concurrent UDF/UDTF (and JAR/SO lookup) fail opening its library. Please mark the entry unavailable to new lookups before waiting (or add a tombstone/retry check under _cache_lock) and coordinate replacement so a retiring entry cannot return a path that cleanup removes.


// 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();
}
RETURN_IF_ERROR(PythonServerManager::instance().clear_module_cache(entry->lib_file));
}

{
std::lock_guard<std::mutex> l(_cache_lock);
auto it = _entry_map.find(fid);
if (it != _entry_map.end() && it->second == entry) {
_entry_map.erase(it);
}
}
// Mark for deletion, destructor will delete the files
entry->should_delete_library.store(true);
return Status::OK();
}

} // namespace doris
2 changes: 1 addition & 1 deletion be/src/runtime/user_function_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ class UserFunctionCache {
std::string* libpath);

// Drop the cached function library by function id.
void drop_function_cache(int64_t fid);
Status drop_function_cache(int64_t fid);

#ifndef BE_TEST
private:
Expand Down
40 changes: 35 additions & 5 deletions be/src/udf/python/python_server.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <butil/fd_utility.h>
#include <dirent.h>
#include <fmt/core.h>
#include <rapidjson/document.h>
#include <signal.h>
#include <sys/poll.h>
#include <sys/stat.h>
Expand Down Expand Up @@ -707,6 +708,7 @@ Status PythonServerManager::_broadcast_action_to_processes(const std::string& ac
int success_count = 0;
int fail_count = 0;
bool has_active_process = false;
std::string failure_details;

for (auto& [version, versioned_pool] : _snapshot_process_pools()) {
std::lock_guard<std::mutex> lock(versioned_pool->mutex);
Expand Down Expand Up @@ -741,12 +743,39 @@ Status PythonServerManager::_broadcast_action_to_processes(const std::string& ac
}

auto result = (*result_stream)->Next();
if (result.ok() && *result) {
success_count++;
} else {
if (!result.ok() || !*result || !(*result)->body) {
fail_count++;
continue;
}

std::string result_body = (*result)->body->ToString();
rapidjson::Document result_json;
result_json.Parse(result_body.data(), result_body.size());
bool action_succeeded = false;
if (!result_json.HasParseError() && result_json.IsObject()) {
auto success = result_json.FindMember("success");
action_succeeded = success != result_json.MemberEnd() &&
success->value.IsBool() && success->value.GetBool();
}
if (!action_succeeded) {
fail_count++;
std::string error = "invalid action result";
if (!result_json.HasParseError() && result_json.IsObject()) {
auto error_member = result_json.FindMember("error");
if (error_member != result_json.MemberEnd() &&
error_member->value.IsString()) {
error.assign(error_member->value.GetString(),
error_member->value.GetStringLength());
}
}
if (!failure_details.empty()) {
failure_details.append("; ");
}
failure_details.append(fmt::format("{}: {}", process->get_uri(), error));
continue;
}

success_count++;
} catch (...) {
fail_count++;
}
Expand All @@ -761,8 +790,9 @@ Status PythonServerManager::_broadcast_action_to_processes(const std::string& ac
<< ", failed=" << fail_count;

if (fail_count > 0) {
return Status::InternalError("{} failed for {}, success={}, failed={}", action_type,
log_name, success_count, fail_count);
return Status::InternalError("{} failed for {}, success={}, failed={}, errors=[{}]",
action_type, log_name, success_count, fail_count,
failure_details);
}

return Status::OK();
Expand Down
Loading
Loading