.. _pipeline-parallelism: Pipeline Parallelism ==================== .. currentmodule:: spdl.pipeline The :py:class:`Pipeline` class supports multi-threading and multi-processing. You can also use a ``Pipeline`` object as the source iterator of another ``Pipeline``. When experimenting, this flexibility makes it easy to switch multi-threading, multi-processing and mixtures of them. This page covers the fundamentals: the executor mechanism, dedicated thread pools, the cost of crossing a process boundary, the :py:meth:`~PipelineBuilder.to` region API, and :py:func:`run_pipeline_in_subprocess`. For the higher-level patterns that combine these primitives into an efficient data loader (MT, MP and MTP), see :ref:`execution-models`. Specifying an executor ---------------------- The core mechanism to deploy concurrency is :py:meth:`asyncio.loop.run_in_executor` method. The synchronous function (or generator) provided to ``Pipeline`` is executed asynchronously using the ``run_in_executor`` method. When you provide a synchronous function (or generator), the ``PipelineBuilder`` internally converts it to an asynchronous equivalent using the ``run_in_executor`` method. In the following snippet, an ``executor`` argument is provided when constructing the ``Pipeline``. .. code-block:: executor: ThreadPoolExecutor | ProcessPoolExecutor | None = ... def my_func(input): ... pipeline = ( PipelineBuilder() .add_source(...) .pipe(my_func, executor=executor) .add_sink(...) .build(...) ) Internally, the ``my_func`` function is converted to an asynchronous equivalent, meaning it's dispatched to the provided executor (or a default one if the executor is ``None``) as follows. .. code-block:: async asynchronous_my_func(input): loop = asyncio.get_running_loop() coroutine = loop.run_in_executor(executor, my_func, input) return await coroutine Multi-threading (default) ------------------------- If you build a pipeline without any customization, it defaults to multi-threading. The event loop dispatches the tasks to the default :py:class:`~concurrent.futures.ThreadPoolExecutor` created with the maximum concurrency specified in :py:meth:`PipelineBuilder.build` method. .. mermaid:: %%{init: {'theme':'base'}}%% graph TB subgraph Process["Process"] EL["Event Loop
(Main Thread)"] subgraph TP["Thread Pool"] T1["Worker Thread 1"] T2["Worker Thread 2"] T3["Worker Thread 3"] T4["Worker Thread 4"] end EL -->|schedules tasks| TP TP -->|returns results| EL end style Process fill:#f0f8ff style EL fill:#e1f5ff style TP fill:#fff4e1 style T1 fill:#fffacd style T2 fill:#fffacd style T3 fill:#fffacd style T4 fill:#fffacd .. note:: **Multi-threading characteristics:** - All threads (main thread and worker threads) run within a single process and naturally share the same memory address space - Fast task startup and minimal overhead - Data can be passed by reference (no copying needed) with fast inter-thread communication - Constrained by the GIL for Python code - best for I/O-bound tasks or GIL-releasing operations **GIL considerations:** To take advantage of multi-threading, your stage functions must mainly consist of operations that release the GIL. Libraries such as PyTorch and NumPy release the GIL when manipulating arrays, and SPDL offers efficient GIL-releasing functions for decoding raw bytes through the :py:mod:`spdl.io` module. See :ref:`working-around-the-gil` for which libraries release the GIL, and what to do when a stage's function does not. .. _pipeline-parallelism-custom-mt: Multi-threading (custom) ------------------------ There are cases where you want to use a dedicated thread for certain task. #. You need to maintain a state across multiple task invocations. (caching for faster execution or storing the application context) #. You want to specify a different number of concurrency. One notable example that meets these conditions is transferring data between the CPU and GPU. CUDA streams provide independent scheduling lanes, but they do not bypass the GPU's copy-engine limits. On the hardware configuration discussed here, at most one host-to-device transfer and one device-to-host transfer can run at the same time. Transfers in opposite directions can overlap each other, while additional transfers in the same direction are serialized. Compute kernels use separate execution resources, so multiple kernels can overlap these transfers and one another when the GPU has sufficient resources. Creating more CUDA streams does not by itself increase the amount of concurrent work. To overlap a transfer with model computation, use a CUDA stream dedicated to that transfer direction and reuse the same stream across function invocations. To maintain a state, you can either encapsulate it in a callable class instance, or put it in a `thread-local storage `_. The following example shows how to initialize and store a CUDA stream in a thread-local storage. .. note:: The following code is now available as :py:func:`spdl.io.transfer_tensor_h2d`. .. code-block:: python import threading THREAD_LOCAL = threading.local() def _get_threadlocal_stream(index: int) -> tuple[torch.cuda.Stream, torch.device]: if not hasattr(THREAD_LOCAL, "stream"): device = torch.device(f"cuda:{index}") THREAD_LOCAL.stream = torch.cuda.Stream(device) THREAD_LOCAL.device = device return THREAD_LOCAL.stream, THREAD_LOCAL.device The following code illustrates a way to transfer data using the same dedicated stream across function invocations. .. code-block:: python def transfer_data(data: list[Tensor], index: int = 0): stream, device = _get_threadlocal_stream(index) with torch.cuda.stream(stream): data = [ t.obj.pin_memory().to(device, non_blocking=True) for t in data] stream.synchronize() return data Now we want to run this function in background, but we want to use only one thread, and keep using the same thread. For this purpose we create a ``ThreadPoolExecutor`` with one thread and pass it to the pipeline. .. code-block:: python transfer_executor = ThreadPoolExecutor(max_workers=1) pipeline = ( PipelineBuilder() .add_source(...) .pipe(...) .pipe(transfer_data, executor=transfer_executor) .add_sink(...) ) This way, the transfer function is always executed in a dedicated thread, so that it keeps using the same CUDA stream. `PyTorch Profiler `_ can verify that the background thread issues transfers on a dedicated CUDA stream while the foreground stream executes model work. The following trace shows host-to-device transfer using this pattern. .. image:: ../../_static/data/parallelism_transfer.png :alt: Host-to-device transfer overlapping foreground CUDA work. The same pattern applies when moving model outputs from the GPU back to the CPU. :py:func:`spdl.io.transfer_tensor_d2h` uses a dedicated CUDA stream and a direct non-blocking transfer into page-locked CPU tensors. Run it on a persistent single-worker executor so that consecutive calls reuse the cached stream while the foreground thread continues to submit independent work to its compute stream. .. code-block:: python d2h_executor = ThreadPoolExecutor(max_workers=1) with d2h_executor: ready = torch.cuda.Event() ready.record() future = d2h_executor.submit( spdl.io.transfer_tensor_d2h, model_output, producer_event=ready, ) # Independent work submitted to the foreground CUDA stream can overlap # with the device-to-host transfer. next_output = model(next_batch) cpu_output = future.result() The event is recorded on the producer thread before dispatch, so the transfer stream waits for exactly the work that produces ``model_output`` without also waiting for later foreground submissions. When :py:func:`~spdl.io.transfer_tensor_d2h` is called directly, it records this event automatically on the calling thread's current stream. See :py:mod:`benchmark_transfer_overlap` for a reproducible benchmark of serialized and concurrent transfer-plus-compute makespans. The following Perfetto view shows the corresponding device-to-host pattern. Four ``Memcpy DtoH`` operations on stream 13 overlap the foreground SGEMM kernels on the default stream (stream 7). In this capture, 19.4 ms of the 21.3 ms device-to-host transfer overlaps compute (91%). .. image:: ../../_static/data/parallelism_transfer_d2h.png :alt: Device-to-host transfer on stream 13 overlapping SGEMM kernels on the default stream. :width: 100% Multi-processing (stage) ------------------------ What if a function does not release the GIL? If a stage relies on a function that takes a long time to execute (for example, a network utility, or a long-running pure-Python routine) but does not release the GIL, running it on the shared thread pool blocks the other stages. You can instead delegate that single stage to a subprocess. :py:meth:`PipelineBuilder.pipe` takes an optional ``executor`` argument. The default behavior is to use the thread pool shared among all stages. If you pass an instance of :py:class:`~concurrent.futures.ProcessPoolExecutor`, that stage executes the function in a subprocess, free of the GIL. .. code-block:: executor = ProcessPoolExecutor(max_workers=num_processes) pipeline = ( PipelineBuilder() .add_source(src) .pipe(stage1, executor=executor, concurrency=num_processes) .pipe(stage2, ...) .pipe(stage3, ...) .add_sink(1) .build() ) This builds a pipeline like the following, where ``stage1`` runs in the process pool while ``stage2`` and ``stage3`` stay on the main process's thread pool. .. include:: ../plots/faq_subprocess_chart.txt .. note:: **Multi-processing characteristics:** - Each process has its own isolated memory space - No GIL constraints - true parallelism for CPU-bound tasks - Data must be pickled and copied between processes (overhead) - Slower startup due to process creation - Best for CPU-bound tasks that hold the GIL .. note:: Along with the function arguments and the return value, the function (callable) itself is serialized and sent to the subprocess, then the result is sent back. Therefore the function, its input argument, and its output value must all be `picklable `_. Plain functions work; closures and class methods cannot be passed. If you want to bind extra arguments to a function, you can use :py:func:`functools.partial`. If you want to pass around an object that's not picklable by default, you can define the serialization protocol by providing :py:meth:`object.__getstate__` and :py:meth:`object.__setstate__`. .. tip:: If you need to perform one-time initialization in the subprocess, you can use the ``initializer`` and ``initargs`` arguments. The values passed as ``initializer`` and ``initargs`` must be picklable. If constructing an object in a process that does not support pickle, then you can pass constructor arguments instead and store the resulting object in global scope. See also https://stackoverflow.com/a/68783184/3670924. Example .. code-block:: def init_resource(*args): global rsc rsc = ResourceClass(*args) def process_with_resource(item): global rsc return rsc.process(item) executor = ProcessPoolExecutor( max_workers=4, mp_context=None, initializer=init_resource, initargs=(...), ) pipeline = ( PipelineBuilder() .add_source() .pipe( process_with_resource, executor=executor, concurrency=4, ) .add_sink(3) .build() ) The cost of crossing a process boundary ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Every time a stage is dispatched to a subprocess, its input is pickled and copied into the worker, and the result is pickled and copied back. For a single heavy, GIL-holding stage this overhead is easily worth it. See :ref:`ipc-cost` for why this pickle-and-copy is expensive and how it scales with the payload. The *size* of what crosses the boundary matters too, since pickling and copying a large payload is costly and is paid on both sides. PyTorch tensors (and NumPy arrays) are the exception: their buffers are moved through shared memory and only a handle is pickled, so transferring them is far cheaper than their size suggests — even though a fresh shared-memory segment is allocated on each pickling. Prefer passing a :py:class:`torch.Tensor` across the boundary. To send a raw byte string cheaply, write it to shared memory yourself, or wrap it in a 1-D ``uint8`` tensor so it rides the same fast path. To squeeze the last bit of IPC performance, :py:class:`~spdl.pipeline.SharedMemorySegmentPool` reuses a pre-allocated pool of shared-memory segments instead of allocating one per transfer, edging out a plain tensor. See the :ref:`shared-memory arena case study ` and the :py:mod:`benchmark_arena_transport` example for the measured numbers. But if you send several consecutive stages each to a ``ProcessPoolExecutor``, the data makes a round trip between the main process and a worker at *every* stage boundary. Copying data back and forth like this wastes both time and memory. To keep several stages running together inside the workers without paying a per-stage round trip, mark a region of the pipeline (next section) or run the whole pipeline in a subprocess. These are the building blocks behind the recommended MP and MTP patterns; see :ref:`execution-models`. Multi-processing (region) ------------------------- Rather than paying a process round trip per stage, you can mark a **region** of the pipeline to run together in a worker pool with :py:meth:`~spdl.pipeline.PipelineBuilder.to`. Every stage between ``.to(ProcessPoolExecutorConfig(...))`` and ``.to(MAIN_PROCESS)`` runs as one nested :py:class:`Pipeline` inside a pool of worker processes: .. code-block:: from spdl.pipeline.defs import MAIN_PROCESS, ProcessPoolExecutorConfig pipeline = ( PipelineBuilder() .add_source(...) .to(ProcessPoolExecutorConfig(max_workers=4)) .pipe(op1, concurrency=2) # runs in a worker process .aggregate(batch_size) # runs in a worker process .pipe(op2, concurrency=3) # runs in a worker process .to(MAIN_PROCESS) # data returns to the main process .add_sink(...) .build(num_threads=...) ) Because the region's stages run back-to-back inside one worker, the value handed from one stage to the next is **not** copied back to the main process between them. This removes the inter-stage IPC entirely, and — unlike the per-stage multi-processing above — those intermediate values do **not** need to be picklable; only the region's inputs and outputs cross the process boundary. Each stage keeps its own ``concurrency`` and its own per-stage performance stats (the nested pipeline is built with the usual hooks, so the stats are reported from inside the worker). .. note:: A stage's ``concurrency`` applies *within each worker process*, so the effective concurrency of a stage across the pool is ``concurrency × max_workers``. For example, ``.pipe(op, concurrency=2)`` in a region with ``max_workers=4`` runs up to 8 invocations of ``op`` at once. Size each stage's ``concurrency`` together with ``max_workers`` to stay within your CPU budget. Unlike passing ``executor=`` to individual :py:meth:`~spdl.pipeline.PipelineBuilder.pipe` calls, a region also carries :py:meth:`~spdl.pipeline.PipelineBuilder.aggregate`, :py:meth:`~spdl.pipeline.PipelineBuilder.disaggregate`, and :py:meth:`~spdl.pipeline.PipelineBuilder.path_variants` stages into the worker, and gives the worker-pool configuration (worker count, ``mp_context``, ``initializer``) a single home. A pipeline starts on the main process, so a region is opened by a ``.to(ProcessPoolExecutorConfig(...))`` and closed by ``.to(MAIN_PROCESS)``; the region must be closed before :py:meth:`~spdl.pipeline.PipelineBuilder.add_sink`. Inside a region, generator ops and async ops work as usual, a :py:meth:`continuous source ` keeps the worker sub-pipelines warm across epochs, and a region composes with :py:func:`~spdl.pipeline.run_pipeline_in_subprocess`. Every op must be picklable, since the region config is shipped to the worker. See the :py:meth:`~spdl.pipeline.PipelineBuilder.to` reference for the details. Buffering the boundary ^^^^^^^^^^^^^^^^^^^^^^ Crossing the region boundary costs a fixed amount *per transfer* — pickling, a queue handoff, and (for tensors) a shared-memory allocation — no matter how small the item is. When the items are small, that fixed cost dominates. Pass ``buffer_size`` to :py:meth:`~PipelineBuilder.to` to buffer several items into one transfer: .. code-block:: pipeline = ( PipelineBuilder() .add_source(rows) # small per-row items .to(ProcessPoolExecutorConfig(max_workers=4), buffer_size=32) .pipe(decode) # still receives ONE row at a time .aggregate(batch_size) # independent of buffer_size .to(MAIN_PROCESS, buffer_size=2) # 2 batches per transfer coming back .add_sink() .build(num_threads=...) ) Each marker sizes the handoff that happens where it sits, so a region's two directions are configured independently: the ``to()`` that opens the region sizes the data flowing into its workers, and the one that closes it sizes the results flowing back. That separation matters because the two directions rarely carry comparable items — in the example above, single rows go in and whole batches come back, so one number could not be right for both. A marker sitting between two adjacent regions sizes that single handoff for both sides. The buffering is transparent: the boundary packs items on one side and unpacks them on the other, so stages on both sides — and every stage after the region — still see individual items. In particular ``buffer_size`` is independent of any :py:meth:`~spdl.pipeline.PipelineBuilder.aggregate` you use, so the training batch size and the IPC transfer size can be tuned separately. The default, ``1``, transfers one item at a time. It is an upper bound rather than a fixed grouping: a partial buffer is flushed at the end of a stream and at every epoch boundary, so the last transfer of each is usually short. Nothing assumes the region emits as many items as it received: a ``None``-returning op, a generator op, and an ``aggregate`` inside the region all work unchanged. The cost is latency and memory. Up to ``buffer_size - 1`` items are held while a transfer fills, so a value well above what the producing side can keep up with delays the first results; and because each queued transfer holds up to ``buffer_size`` items, the in-flight item bound scales with it. Start at the batch size you would otherwise have aggregated before the region, then tune. .. note:: Placing ``aggregate`` *inside* a region means each worker aggregates only the items routed to it, so batches are formed per worker rather than globally. With ``drop_last=True`` that discards up to ``max_workers × (batch_size - 1)`` items per epoch instead of ``batch_size - 1``. To run a region in **subinterpreters** (Python 3.14+) instead of subprocesses, pass a :py:class:`~spdl.pipeline.defs.InterpreterPoolExecutorConfig`; the region's ops must avoid NumPy/PyTorch, which cannot be imported in a subinterpreter. .. note:: A region produces results in completion order across its pool workers, so a stage built with ``output_order="input"`` cannot appear inside a region. Running a pipeline in a subprocess ---------------------------------- The :py:func:`spdl.pipeline.run_pipeline_in_subprocess` function moves a given :py:class:`PipelineBuilder` configuration to a subprocess, builds and runs the :py:class:`Pipeline` there, and delivers the results to the main process through an inter-process queue. .. mermaid:: %%{init: {'theme':'base'}}%% graph TB subgraph SP["Subprocess"] SEL["Event Loop
(Sub Thread)"] subgraph STP["Thread Pool"] ST1["Stage 1"] ST2["Stage 2"] ST3["Stage 3"] end SEL -->|schedules| STP end subgraph MP["Main Process"] MEL["Iterates results
(Main Thread)"] end STP -->|results| Q["IPC Queue"] Q -->|results| MEL style MP fill:#e1f5ff style SP fill:#ffe1e1 style MEL fill:#b3d9ff style SEL fill:#ffb3b3 style STP fill:#fff4e1 style Q fill:#e8f5e9 style ST1 fill:#fffacd style ST2 fill:#fffacd style ST3 fill:#fffacd .. note:: **How it works:** - **Subprocess**: runs the full pipeline with its own event loop and thread pool. - **IPC Queue**: results are transferred to the main process via inter-process communication. - **Main process**: iterates the returned iterable to receive the results. The following example shows how to use the function. .. code-block:: python # Construct a builder and get its config builder = ( spdl.pipeline.PipelineBuilder() .add_source(...) .pipe(...) ... .add_sink(...) ) config = builder.get_config() # Move it to the subprocess, build the Pipeline iterable = run_pipeline_in_subprocess(config, num_threads=...) # Iterate - epoch 0 for item in iterable: ... # Iterate - epoch 1 for item in iterable: ... .. note:: **Advanced Usage:** - **Pipelines with Merge**: You can run pipelines with sub-pipelines constructed using ``Merge`` by directly passing the ``PipelineConfig`` object to ``run_pipeline_in_subprocess``. This allows complex pipeline topologies to be executed in a subprocess. - **Subinterpreter Execution**: For Python 3.14 and above, the :py:func:`run_pipeline_in_subinterpreter` function is also available. It executes the pipeline in a separate interpreter within the same process, providing interpreter-level isolation while being lighter weight than a full subprocess. Because the result of ``run_pipeline_in_subprocess`` is an ``iterable``, you can build another :py:class:`Pipeline` on top of it. See :ref:`execution-models` for the execution patterns that build on this primitive.