spdl.dataloader.PipelineBuilder

class PipelineBuilder

Build Pipeline object.

See Pipeline for details.

Methods

add_sink(buffer_size)	Attach a buffer to the end of the pipeline.
add_source(source, **kwargs)	Attach an iterator to the source buffer.
aggregate(num_items, /, *[, drop_last, ...])	Buffer the items in the pipeline.
build(*[, num_threads])	Build the pipeline.
disaggregate(*[, hooks, report_stats_interval])	Disaggregate the items in the pipeline.
pipe(op, /, *[, concurrency, executor, ...])	Apply an operation to items in the pipeline.

add_sink(buffer_size: int) → PipelineBuilder

Attach a buffer to the end of the pipeline.

 │
┌▼┐
│ │ buffer
└─┘

Parameters:

buffer_size – The size of the buffer. Pass 0 for unlimited buffering.

add_source(source: Iterable[T] | AsyncIterable[T], **kwargs) → PipelineBuilder

Attach an iterator to the source buffer.

┌─────────────────┐
│ (Async)Iterator │
└───────┬─────────┘
        ▼

Parameters:

source –

A lightweight iterator that generates data.

Warning

The source iterator must be lightweight as it is executed in async event loop. If the iterator performs a blocking operation, the entire pipeline will be blocked.

aggregate(num_items: int, /, *, drop_last: bool = False, hooks: Sequence[PipelineHook] | None = None, report_stats_interval: float | None = None) → PipelineBuilder

Buffer the items in the pipeline.

Parameters:

• num_items – The number of items to buffer.

• drop_last – Drop the last aggregation if it has less than num_aggregate items.

• hooks – See pipe().

• report_stats_interval – See pipe().

build(*, num_threads: int | None = None) → Pipeline

Build the pipeline.

Parameters:

num_threads – The number of threads in the thread pool attached to async event loop. If not specified, the maximum concurrency value is used.

disaggregate(*, hooks: Sequence[PipelineHook] | None = None, report_stats_interval: float | None = None) → PipelineBuilder

Disaggregate the items in the pipeline.

Parameters:

• hooks – See pipe().

• report_stats_interval – See pipe().

pipe(op: Callable[[T], U] | Callable[[T], Iterable[U]] | Callable[[T], Awaitable[U]] | Callable[[T], AsyncIterable[U]], /, *, concurrency: int = 1, executor: type[Executor] | None = None, name: str | None = None, hooks: Sequence[PipelineHook] | None = None, report_stats_interval: float | None = None, output_order: str = 'completion', **kwargs) → PipelineBuilder

Apply an operation to items in the pipeline.

      │
┌─────▼─────┐
│ Operation │
└─────┬─────┘
      ▼

Parameters:

• op –

  A function applied to items in the queue. The function must take exactly one argument, which is the output from the upstream. If passing around multiple objects, take them as a tuple or use dataclass and define a custom protocol.

  Optionally, the op can be a generator function, async function or async generator function.

  If op is (async) generator, the items yielded are put in the output queue separately.

  Warning

  If op is synchronous geneartor, and executor is an instance of concurrent.futures.ProcessPoolExecutor, the output items are not put in the output queue until the generator is exhausted.

  Async generator, or synchronous generator without ProcessPoolExecutor does not have this issue, and the yielded items are put in the output queue immediately.

  Tip

  When passing an async op, make sure that the op does not call sync function inside. If calling a sync function, use asyncio.loop.run_in_executor() or asyncio.to_thread() to delegate the execution to the thread pool.

• concurrency – The maximum number of async tasks executed concurrently.

• executor –

  A custom executor object to be used to convert the synchronous operation into asynchronous one. If None, the default executor is used.

  It is invalid to provide this argument when the given op is already async.

• name – The name (prefix) to give to the task.

• hooks – Hook objects to be attached to the stage. Hooks are intended for collecting stats of the stage. If None, a default hook, TaskStatsHook is used.

• report_stats_interval – The interval (in seconds) to log the stats of this stage when no hooks is provided. This argument is passed to TaskStatsHook. This argument is effective only when hooks are not provided. If hooks is provided and stats report is needed, the hooks argument should include one of TaskStatsHook instance with the desired interval.

• output_order – If "completion" (default), the items are put to output queue in the order their process is completed. If "input", then the items are put to output queue in the order given in the input queue.