#826 [wip] Add InferencePipeline: 3-stage async pipeline for batch inference
Pull request opened by
kumare3
InferencePipeline wires together preprocessing, batched GPU inference (via DynamicBatcher), and postprocessing with proper thread pool dispatch and backpressure — so users only define three functions:
preprocess_fn (per-item, sync or async, optional executor)
inference_fn (batched, async — the DynamicBatcher process_fn)
postprocess_fn (per-item, receives original input + result)
Key design decisions:
• No new abstraction for batching — reuses DynamicBatcher internally
• Bounded asyncio.Queue between preprocess and batcher for backpressure
• Supports both sync and async stage functions (auto-detected)
• Sync stages can be dispatched to a ThreadPoolExecutor
• postprocess_fn receives the original raw item alongside the inference result so it has full context (filename, metadata, etc.)
• Streaming via async generator (run) or batch collect (run_all)
• Exposes underlying batcher + stats for monitoring/tuning
This pattern emerged from the embed_wikipedia and multimodal-retrieval examples, where every GPU inference use case independently wired the same load → preprocess → batch-infer → postprocess pipeline. The abstraction eliminates that boilerplate while keeping full control over DynamicBatcher tuning knobs (cost budgeting, batch size, timeouts).
flyteorg/flyte-sdk