MPIRE (MultiProcessing Is Really Easy)

repository·master·Indexed 24 days ago

https://github.com/sybrenjansen/mpire

A high-performance Python multiprocessing library designed as a faster, more user-friendly alternative to the standard multiprocessing module. It features a WorkerPool as a drop-in replacement for multiprocessing.Pool, worker state management to avoid reloading heavy resources, copy-on-write shared objects for Linux/macOS, and built-in progress bars via tqdm. Additional tools include a monitoring dashboard, worker insights for profiling overhead, and support for advanced serialization using dill.

Tokens
19.6K
Snippets
58
Records
96
Agent score
84%

What's inside MPIRE

  1. Overview of MPIRE features

    master

    MPIRE (MultiProcessing Is Really Easy) is a Python package designed to be faster and more user-friendly than the standard multiprocessing module. It combines the map-like functionality of multiprocessing.Pool with the benefits of copy-on-write shared objects (available on Linux/macOS via the fork start method).

    Key capabilities include:

    • Map-like functions: Supports map, map_unordered, imap, imap_unordered, apply, and apply_async.
    • Worker State Management: Each worker can maintain its own state using worker_init and worker_exit functions. This is ideal for loading large models or datasets once per worker to avoid expensive serialization.
    • Progress Tracking: Built-in support for tqdm progress bars (including rich and notebook widgets) and progress dashboards.
    • Efficiency & Control: Automatic task chunking, adjustable maximum active tasks to manage memory, and automatic worker restarting to reduce memory footprint.
    • Advanced Execution: CPU pinning, nested pools (via the daemon option), and support for exotic objects (lambdas, etc.) when using dill via multiprocess_.
    • Robustness: Graceful exception handling and support for timeouts on tasks and worker lifecycle functions.
  2. Caveats of using keep_alive with worker_init and worker_exit

    master

    When keep_alive is enabled, workers are not restarted between map calls. This has significant implications for worker_init and worker_exit functions:

    • worker_init: Only executes when a worker is first started. If you change the worker_init function in a subsequent map call, the new function will not be called for existing workers.
    • worker_exit: Only executes when a worker is terminated. If you change the worker_exit function in a subsequent map call, the old workers will not execute the new exit function until they are eventually restarted or the pool is shut down.

    Note on worker_lifespan: If worker_lifespan is enabled, some workers might reach their lifespan limit and restart during a map call. In this case, the restart will trigger the worker_exit function of the worker being replaced, but the timing and which function is called can become complex if functions are changed between calls.

  3. How worker_lifespan affects exit results

    master

    If you use the worker_lifespan option to restart workers during execution, the worker_exit function will be called for every worker that shuts down, and worker_init will be called for every new worker created.

    Because workers may be restarted multiple times, the number of elements returned by pool.get_exit_results() may be greater than the initial n_jobs count.

  4. Manage worker state with worker_init and worker_exit

    master

    You can use worker_init to load resources (like models or database connections) into a worker-specific state once, rather than passing them with every task. To enable this, set use_worker_state=True in the WorkerPool constructor. You can also use worker_exit to run cleanup code when a worker terminates.

    def init(worker_state):
        # Load a big dataset or model and store it in a worker specific worker_state
        worker_state['dataset'] = ...
        worker_state['model'] = ...
    
    def task(worker_state, idx):
        # Let the model predict a specific instance of the dataset
        return worker_state['model'].predict(worker_state['dataset'][idx])
    
    with WorkerPool(n_jobs=5, use_worker_state=True) as pool:
        results = pool.map(task, range(10), worker_init=init)
  5. Understand the structure of worker insights

    master

    The insights dictionary provides detailed metrics per worker and aggregate statistics. For a pool with $N$ workers, containers like n_completed_tasks, start_up_time, init_time, waiting_time, working_time, and exit_time will contain $N$ entries.

    Key Metrics:

    • n_completed_tasks: Number of tasks completed by each worker.
    • start_up_time: Time taken for worker startup.
    • init_time: Time spent in the worker_init function.
    • waiting_time: Time spent waiting for new tasks.
    • working_time: Time spent executing the main task function.
    • exit_time: Time spent in the worker_exit function.
    • Ratios: The dictionary includes mean, std (standard deviation), and ratio (the specific metric divided by the total time) for each category. A higher working_ratio generally indicates a more efficient multiprocessing setup.

    Longest Running Tasks:

    Insights include tracking for the top 5 longest-running tasks:

    • top_5_max_task_durations: A list of the 5 longest durations (sorted descending).
    • top_5_max_task_args: A list of the arguments passed to those specific tasks. The index in this list corresponds to the index in the duration list.
    # Example output structure
    {
     'n_completed_tasks': [28, 24, 24, 24],
     'total_start_up_time': '0:00:00.038',
     'total_init_time': '0:00:00',
     'total_waiting_time': '0:00:00.798',
     'total_working_time': '0:00:04.980',
     'total_exit_time': '0:00:00',
     'total_time': '0:00:05.816',
     'start_up_time': ['0:00:00.010', '0:00:00.008', '0:00:00.008', '0:00:00.011'],
     'start_up_time_mean': '0:00:00.009',
     'start_up_time_std': '0:00:00.001',
     'start_up_ratio': 0.006610452621805033,
     'init_time': ['0:00:00', '0:00:00', '0:00:00', '0:00:00'],
     'init_time_mean': '0:00:00',
     'init_time_std': '0:00:00',
     'init_ratio': 0.0,
     'waiting_time': ['0:00:00.309', '0:00:00.311', '0:00:00.165', '0:00:00.012'],
     'waiting_time_mean': '0:00:00.199',
     'waiting_time_std': '0:00:00.123',
     'waiting_ratio': 0.13722942739284952,
     'working_time': ['0:00:01.142', '0:00:01.135', '0:00:01.278', '0:00:01.423'],
     'working_time_mean': '0:00:01.245',
     'working_time_std': '0:00:01.117',
     'working_ratio': 0.8561601182661567,
     'exit_time': ['0:00:00', '0:00:00', '0:00:00', '0:00:00'],
     'exit_time_mean': '0:00:00',
     'exit_time_std': '0:00:00',
     'exit_ratio': 0.0,
     'top_5_max_task_durations': ['0:00:00.099', '0:00:00.098', '0:00:00.097', '0:00:00.096', '0:00:00.095'],
     'top_5_max_task_args': ['Arg 0: 99', 'Arg 0: 98', 'Arg 0: 97', 'Arg 0: 96', 'Arg 0: 95']
    }
  6. How to use nested WorkerPools

    master

    By default, WorkerPool spawns daemon child processes, which prevents them from creating their own child processes (nested pools). To enable nested pools, you must set daemon=False.

    Important Considerations for Nested Pools:

    • Start Method: It is highly recommended to use start_method='spawn' to avoid thread-safety issues when spawning processes.
    • Threading: Nested pools are not supported when using threading.
    • Forkserver: If using forkserver in a nested pool, the outer pool must also use either spawn or forkserver due to Python limitations.
    • Stability: Nested pools are not considered production-ready; they may occasionally cause deadlocks during error handling or keyboard interrupts.
    def job(...):
        with WorkerPool(n_jobs=4) as p:
            # Do some work
            results = p.map(...)
    
    # This will raise an AssertionError because daemon=True
    with WorkerPool(n_jobs=4, daemon=True, start_method='spawn') as pool:
        pool.map(job, ...)
    
    # This will work because daemon=False
    with WorkerPool(n_jobs=4, daemon=False, start_method='spawn') as pool:
        pool.map(job, ...)
  7. How shared objects work in MPIRE

    master

    MPIRE allows you to provide shared objects to workers, similar to multiprocessing.Process. The behavior of these objects depends on the start_method used:

    • fork: Shared objects are treated as copy-on-write. They share the same memory address and are only copied when changes are made. This is ideal for large datasets to avoid memory exhaustion. Note: fork is not available on Windows.
    • threading: Shared objects are readable and writable without any copies being made.
    • spawn and forkserver: Shared objects are copied once for each worker (unlike a regular multiprocessing.Pool which copies for each task).

    Important Implementation Detail: Shared objects are passed as the second argument to your task function, immediately following the worker ID (if worker IDs are enabled).

  8. How task chunking works in MPIRE

    master

    By default, MPIRE chunks tasks into 64 * n_jobs chunks. Each worker receives one chunk of tasks at a time before returning results.

    Chunking is beneficial when tasks are computationally small, as it reduces the overhead of pickling and unpickling data sent between the worker and the main process.

    To calculate chunk sizes, MPIRE needs to know the total number of tasks. This is automatically handled for containers that implement __len__ (like list or tuple), but for generators, you must provide the length manually using the iterable_len parameter to avoid performance degradation.

  9. Use copy-on-write shared objects

    master

    To share large objects (like datasets or models) across all workers without the overhead of serialization/copying for every task, use the shared_objects parameter.

    Note: Copy-on-write is only available when using the fork start method. For threading, objects are shared as-is. For other methods, objects are copied once per worker.

    def time_consuming_function(some_object, x):
        import time
        time.sleep(1)
        return x
    
    def main():
        some_object = ... # Large object
        with WorkerPool(n_jobs=5, shared_objects=some_object) as pool:
            results = pool.map(time_consuming_function, range(10), progress_bar=True)
  10. Understand multiprocessing start methods in MPIRE

    master

    MPIRE supports several multiprocessing start methods via the start_method parameter in WorkerPool. The choice of method affects how child processes are created and how memory is shared:

    • fork: Copies the parent process, making the child effectively identical. It includes everything currently in memory and enables copy-on-write shared objects. This is the default on Unix.
    • spawn: Starts a fresh Python interpreter, inheriting only necessary resources. This is the default on Windows.
    • forkserver: Starts a server process (using spawn) which then forks new processes upon request.
    • threading: Starts child threads instead of processes. This is subject to the Global Interpreter Lock (GIL) but is suitable for I/O-intensive tasks.
    Start methodAvailable on UnixAvailable on Windows
    forkYes (default)No
    spawnYesYes (default)
    forkserverYesNo
    threadingYesYes
  11. Share objects between workers using copy-on-write

    master

    If you have large objects (like datasets or models) that you want to share across all workers without the overhead of serialization or copying, use the shared_objects parameter in WorkerPool.

    Important Requirements:

    • Platform: Copy-on-write is not available on Windows because it requires the fork start method.
    • Usage: Pass the object to shared_objects. MPIRE will pass it to workers only once. The object is only copied for a specific worker if it is modified within that worker's function.
    def time_consuming_function(some_object, x):
        import time
        time.sleep(1)
        return x
    
    def main():
        some_object = ... # Large object to share
        with WorkerPool(n_jobs=5, shared_objects=some_object, start_method='fork') as pool:
            results = pool.map(time_consuming_function, range(10), progress_bar=True)