Python中joblib Parallel内存泄漏定位及多进程线程超时方案问询

Great question—let's break down exactly why your decorator is causing memory bloat and crashes:

Every time get_output runs, you create a new multiprocessing.pool.ThreadPool(1) but never clean it up. Thread pools in Python don't automatically shut down their worker threads when the pool object goes out of scope. Those threads stay alive indefinitely, waiting for new tasks to process.

Since you're using joblib's Parallel with 50 threading workers, each worker thread spawns another thread pool via your decorator. Over time, you end up with hundreds of orphaned threads hogging memory until your environment can't handle it anymore.

The "2x timeout" delay makes sense too: when the timeout triggers, async_result.get(timeout) throws an error, but the thread running do_stuff is still chugging along in the background. Your joblib worker thread waits for that background thread to finish (even after the timeout), leading to longer-than-expected wait times.

The simplest fix is to ensure your thread pool is properly closed after use. Use a with statement to manage the pool—this guarantees it gets shut down automatically, even if an error occurs:

import functools
import multiprocessing.pool

def with_timeout(timeout):
    def decorator(decorated):
        @functools.wraps(decorated)
        def inner(*args, **kwargs):
            # Use context manager to auto-close the pool after use
            with multiprocessing.pool.ThreadPool(1) as pool:
                async_result = pool.apply_async(decorated, args, kwargs)
                try:
                    return async_result.get(timeout)
                except multiprocessing.TimeoutError:
                    return None  # Explicit return value for timeout cases
        return inner
    return decorator

Since you're already using joblib's threading backend, nesting another thread pool is unnecessary. You can implement timeout directly with the threading module, which is lighter and avoids nested thread overhead:

import functools
import threading

def with_timeout(timeout):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            result_container = [None]
            exception_container = [None]
            completion_event = threading.Event()

            def task_runner():
                try:
                    # Execute the target function and store the result
                    result_container[0] = func(*args, **kwargs)
                except Exception as e:
                    # Capture exceptions to re-raise later
                    exception_container[0] = e
                finally:
                    # Signal the task is complete (success or failure)
                    completion_event.set()

            # Use a daemon thread—dies automatically if the main process exits
            task_thread = threading.Thread(target=task_runner)
            task_thread.daemon = True
            task_thread.start()

            # Wait for the task to finish or hit the timeout
            if not completion_event.wait(timeout):
                # Timeout occurred—return None or handle as needed
                return None

            # Re-raise any exceptions caught during task execution
            if exception_container[0] is not None:
                raise exception_container[0]

            return result_container[0]
        return wrapper
    return decorator

Important note: Python doesn't support safe forced termination of threads. If a timeout triggers, the task thread will still run until it finishes on its own. This is a limitation of Python's GIL, but it's safer than trying to kill threads (which can leave resources in inconsistent states).

You mentioned Parallel hangs when len(list) <= n_jobs. While the timeout fix addresses the memory leak, the hang might stem from blocking operations in do_stuff that don't play well with threading (e.g., unhandled locks, blocking IO that doesn't release the GIL). If the hang persists, try:

• Auditing do_stuff for deadlocks or long-running blocking calls
• Testing joblib's loky backend instead of threading (it's the default for CPU-bound tasks, but can also handle IO-bound work more reliably in some cases)

内容的提问来源于stack exchange，提问作者sudonym