Skip to content

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

Draft
linrrzqqq wants to merge 3 commits into
apache:masterfrom
linrrzqqq:pyudf-cache
Draft

[Fix](python udf) isolate module caches across functions#67511
linrrzqqq wants to merge 3 commits into
apache:masterfrom
linrrzqqq:pyudf-cache

Conversation

@linrrzqqq

Copy link
Copy Markdown
Collaborator

What problem does this PR solve?

Fixes Python UDF module isolation and cache cleanup issues when different functions contain modules with the same name.

Python UDF/UDAF/UDTF functions run in shared Python server processes and therefore share process-wide import state such as sys.modules and sys.path. When two function packages contain dependencies with the same module name, one function may reuse the module loaded from another package. This is especially visible with:

  • Imports executed lazily inside function bodies.
  • Imports executed from user-created threads.
  • UDAF state serialization and deserialization.
  • Concurrent module loading.
  • Function cache cleanup while a module is being imported.

Cache cleanup failures could also be ignored by the BE, allowing module files to be deleted while Python processes still referenced them.

What is changed?

  • Isolate imported modules by UDF package location using a per-location module context.
  • Restore the corresponding module context while executing module UDF/UDAF/UDTF code.
  • Intercept __import__ and importlib.import_module so delayed imports use the correct function package.
  • Propagate module context to user-created Python threads.
  • Coordinate normal imports and temporary sys.path/sys.modules modifications with a reader-writer locking scheme.
  • Avoid the global write lock and full sys.modules scan for cached imports.
  • Resolve UDAF state classes directly from the corresponding module context during deserialization.
  • Use the new module import lock when clearing Python module caches.
  • Propagate Python cache cleanup failures to UserFunctionCache.
  • Preserve cache entries and module files when cleanup fails, allowing the cleanup operation to be retried safely.

@hello-stephen

Copy link
Copy Markdown
Contributor

Thank you for your contribution to Apache Doris.
Don't know what should be done next? See How to process your PR.

Please clearly describe your PR:

  1. What problem was fixed (it's best to include specific error reporting information). How it was fixed.
  2. Which behaviors were modified. What was the previous behavior, what is it now, why was it modified, and what possible impacts might there be.
  3. What features were added. Why was this function added?
  4. Which code was refactored and why was this part of the code refactored?
  5. Which functions were optimized and what is the difference before and after the optimization?

Python module cache entries use location strings as keys, but cleanup treated each key as a tuple. Match the full location key and obtain the module name from the cached module before eviction.
@linrrzqqq

Copy link
Copy Markdown
Collaborator Author

/review

@github-actions github-actions Bot left a comment

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.

Automated review found five correctness issues in cache retirement and Python module isolation; details are inline. No builds or test suites were run, per the review instructions.


// 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);

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.

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()) {

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.

return cls._import_locks[cache_key]
@classmethod
@contextmanager
def _shared_module_import(cls):

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.

temporarily_restore_udf_modules overlays udf_modules but leaves unrelated names already present in sys.modules. If an entry module performs a later import dep for a dependency that was not imported during initial load, the wrapper misses the UDF cache and the original importer returns the server's preloaded sys.modules['dep'] without consulting this UDF's sys.path. This silently mixes environments (or raises on a missing attribute). Please mask/save conflicting target names during a UDF miss (or use location-qualified module namespaces) and add a lazy absolute-dependency collision test.

cls._module_import_condition.notify_all()

@classmethod
@contextmanager

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.

use_module_context only binds a ContextVar and does not retain a reader/borrower lease. _handle_exchange_udf applies it separately around each batch, so a paused stream can have its module cache evicted by clear_module_cache and its directory deleted by drop_function_cache between batches. If the next batch lazily imports a dependency, the wrapper restores a path that no longer exists and the running query fails. Please retain a module-context/read lease for the exchange lifetime (or defer eviction/deletion until the stream closes).

)
_current_module_import_operation: contextvars.ContextVar[
Optional[_ModuleImportOperation]
] = contextvars.ContextVar("current_module_import_operation", default=None)

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.

The child-context capture relies on thread_target is builtins.__import__ or thread_target is importlib.import_module. A valid UDF pattern such as Thread(target=functools.partial(importlib.import_module, '.dep', package)) therefore starts with no direct_import_context; the partial invokes the wrapper without UDF caller globals and the dependency is searched only in the server environment, producing a wrong module or ModuleNotFoundError. Please propagate the owning context to adapted callables (or wrap the callable invocation) rather than relying on identity checks.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants