Pipeline Parallelism¶
The Pipeline class supports multi-threading and multi-processing.
You can also use a Pipeline objects as source iterator of another Pipeline.
When experimenting, this flexibility makes it easy to switch multi-threading,
multi-processing and mixtures of them.
Specifying an executor¶
The core mechanism to deploy concurrency is 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.
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.
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
ThreadPoolExecutor
created with the maximum concurrency specified in PipelineBuilder.build()
method.
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:
Before your application can take advantage of free-threaded Python, to properly achieve concurrency, your stage functions must mainly consist of operations that release the GIL.
Libraries such as PyTorch and NumPy release the GIL when manipulating arrays, so they are usually fine.
For loading raw byte strings into array format, SPDL offers efficient
functions through spdl.io module.
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 to the GPU. Due to the hardware constraints, only one data transfer can be performed at a time. To transfer data without interrupting the model training, you need to use a stream object dedicated for the transfer, and you want to keep using the same stream object across multiple 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 spdl.io.transfer_tensor().
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.
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.
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.
When tracing this pipeline with PyTorch Profiler, we can see that it is always the one background thread that issues data transfer, and the transfer overlaps with the stream executing the model training.
Multi-processing (stage)¶
Similar to the custom multi-threading, by providing an instance of
ProcessPoolExecutor, that stage is executed in a
subprocess.
executor = ProcessPoolExecutor(...)
pipeline = (
PipelineBuilder()
.add_source(...)
.pipe(task_function, executor=executor)
.add_sink(...)
)
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 that when you dispatch the stage to subprocess, both the function (callable) and the argument are sent from the main process to the subprocess. Then the result obtained by passing the argument to the function is sent back from the subprocess to the main process. Therefore, all of the function (callable), the input argument and the output value must be picklable.
If you want to bind extra arguments to a function, you can use
functools.partial().
If you want to pass around an object that’s not picklable by default,
you can define the serialization protocol by providing
object.__getstate__() and object.__setstate__().
Multi-processing (combined)¶
If you have multiple stages that you want to run in subprocess, it is inefficient to copy data between processes back and forth.
One workaround is to combine stages and let each process run processes in a batch.
def preprocess(items: list[T]) -> U:
# performs decode/preprocess and collation
...
executor = ProcessPoolExecutor(...)
pipeline = (
PipelineBuilder()
.add_source(...)
.aggregate(batch_size)
.pipe(preprocess, executor=executor, concurrency=...)
.add_sink(...)
)
This approach is similar to the conventional DataLoader. One downside with this approach is less robust in error handling than the previous approaches. If preprocessing fails for one item, and if you want to ensure the size of the batch to be consistent, then all items must be dropped too. The other approach does not suffer from this.
It is also cumbersome: you must hand-combine the stages into one function,
which collapses the per-stage performance stats into a single number and
discards each stage’s individual concurrency.
Multi-processing (fused)¶
Instead of hand-combining stages, you can ask the builder to fuse them for
you by passing fuse_subprocess_stages=True to
build(). Runs of consecutive
pipe() stages that share the same
process-pool (or interpreter-pool) executor instance are fused into a single
stage that runs the run as one nested Pipeline inside the worker
pool:
executor = ProcessPoolExecutor(max_workers=4)
pipeline = (
PipelineBuilder()
.add_source(...)
.pipe(op1, executor=executor, concurrency=2)
.pipe(op2, executor=executor, concurrency=3)
.add_sink(...)
.build(num_threads=..., fuse_subprocess_stages=True)
)
Because op1 and op2 now run back-to-back inside one worker, the
intermediate value 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 — the value handed from op1 to op2 does not
need to be picklable. Each fused 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).
A generator op (a function that yields) is supported as a fused
process-pool stage: each input item fans out into the values the generator
yields, exactly as in an unfused pipeline. As with any sync generator on a
process-pool executor, the yielded items are materialized once the generator is
exhausted rather than streamed out incrementally.
An async op (an async def function or an async generator) can be fused
too. Because an async op always runs on the event loop, it takes no executor to
run it; instead, tag it with the same pool executor as its neighbours and
it joins their fused run, executing on the worker’s own event loop:
.pipe(sync_op, executor=executor)
.pipe(async_op, executor=executor) # runs on the worker's event loop
.pipe(sync_op, executor=executor)
All three fuse into one subprocess run, so an async op between two pool stages no longer splits the run in two. The executor is used only to group the stage, not to run the coroutine; a fused async op must be picklable, like any fused stage. Passing a non-isolating executor (e.g. a thread pool) to an async op is an error. Unfused, the tag is ignored and the async op runs in the main process.
Only adjacent pool stages on the same executor are fused. An
aggregate() or
disaggregate() between two pool stages
is not fused — it runs in the main process and keeps its usual batching
semantics, and it splits the surrounding pool stages into separate runs (so
each side fuses on its own only if it has two or more adjacent pool stages).
The same option is accepted by
spdl.pipeline.run_pipeline_in_subprocess(), where the fused worker
pool is owned by the main process and the run executes in those workers — so
the per-stage round-trip between the pipeline subprocess and the pool is
removed as well.
Fusion also works with a continuous source
(~spdl.pipeline.PipelineBuilder.add_source(..., continuous=True)). The worker sub-pipelines run in
continuous mode and stay warm across epochs, and epoch boundaries are
propagated across the pool: each fused stage emits one epoch boundary per epoch
just like an unfused pipeline.
Note
Fusion preserves results but produces them in completion order across the
pool workers. Stages built with output_order="input" are not fused.
Multi-threading in subprocess¶
The multi-threading in subprocess is a paradigm we found effective in the case study Parallelism and Performance.
The spdl.pipeline.run_pipeline_in_subprocess() function moves the given
instance of PipelineBuilder to a subprocess, build and execute the
Pipeline and put the results to inter-process queue.
Note
How it works:
Subprocess: Runs a full pipeline with its own event loop and thread pool
Data processing: Download, decode, and preprocessing happen in the subprocess
IPC Queue: Batched data is transferred to the main process via inter-process communication
Main Process: Receives batched data and performs GPU transfer in a dedicated thread
Benefit: Separates data loading from GPU operations, reducing main thread overhead
The following example shows how to use the function.
# 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
Mergeby directly passing thePipelineConfigobject torun_pipeline_in_subprocess. This allows complex pipeline topologies to be executed in a subprocess.Subinterpreter Execution: For Python 3.14 and above, the
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.
Since the result of the run_pipeline_in_subprocess is an iterable,
you can build a pipeline on top of it.
This allows to build a pipeline that creates a batch object in a subprocess, then transfer the batch to the GPU in a background thread in the main process. We refer this pattern as MTP (“multi-threading in subprocess”).
# Pipeline that fetches data, loads, then collates.
builder = (
PipelineBuilder()
.add_source(Dataset(...))
.pipe(download, concurrency=...)
.pipe(load, concurrency=...)
.aggregate(batch_size)
.pipe(collate)
.add_sink(...)
)
config = builder.get_config()
src = run_pipeline_in_subprocess(config, num_threads=...)
# Build another pipeline on top of it, which transfers the data to a
# GPU
pipeline = (
PipelineBuilder()
.add_source(src)
.pipe(gpu_transfer)
.add_sink(...)
.build(...)
)
# Iterate
for batch in pipeline:
...
The MTP mode helps the OS to schedule GPU kernel launches from the main thread (where the training loop is running) in timely manner, and reduces the number of Python objects that the Python interpreter in the main process has to handle.
For multi-epoch training, build the subprocess source with a continuous
source (add_source(..., continuous=True)). The pipeline is then built and
started once inside the subprocess and reused across epochs, keeping the
prefetch buffer warm and avoiding a per-epoch rebuild gap. This is especially
important when the training step is short (e.g. small models), where a rebuild
gap would otherwise stall the GPU between epochs. The outer (main-process)
pipeline can likewise use continuous=True so that the GPU-transfer stage
keeps running across epoch boundaries.