Function batch_run

quantmind/flows/batch.py:51–136

Run ``flow_fn`` over ``inputs`` with bounded concurrency. Args: flow_fn: Any flow function with signature ``(input, *, cfg, **kwargs) -> Awaitable[OutputT]``. inputs: Inputs to fan out over. Empty list returns an empty ``BatchResult`` immediately.

(
    flow_fn: Callable[..., Awaitable[OutputT]],
    inputs: list[InputT],
    *,
    cfg: BaseFlowCfg | None = None,
    concurrency: int = 4,
    on_error: Literal["raise", "skip"] = "skip",
    on_progress: Callable[[int, int], None] | None = None,
    **flow_kwargs: Any,
)

async def batch_run(
    flow_fn: Callable[..., Awaitable[OutputT]],
    inputs: list[InputT],
    *,
    cfg: BaseFlowCfg | None = None,
    concurrency: int = 4,
    on_error: Literal["raise", "skip"] = "skip",
    on_progress: Callable[[int, int], None] | None = None,
    **flow_kwargs: Any,
) -> BatchResult[OutputT]:
    """Run ``flow_fn`` over ``inputs`` with bounded concurrency.

    Args:
        flow_fn: Any flow function with signature
            ``(input, *, cfg, **kwargs) -> Awaitable[OutputT]``.
        inputs: Inputs to fan out over. Empty list returns an empty
            ``BatchResult`` immediately.
        cfg: Shared cfg forwarded to every call. ``None`` lets the flow
            use its own default.
        concurrency: Maximum number of in-flight calls. Must be ≥ 1.
        on_error: ``"raise"`` propagates the first failure (siblings get
            cancelled); ``"skip"`` records every failure into
            ``errors`` and returns the batch normally.
        on_progress: Called as ``on_progress(done, total)`` after every
            completion (success or failure). Must be cheap and
            non-blocking — callbacks are invoked synchronously inside
            the worker loop.
        **flow_kwargs: Forwarded verbatim to ``flow_fn``. ``memory=`` is
            **forbidden** in MVP; passing it raises ``ValueError``.

    Returns:
        ``BatchResult`` with ``results`` parallel to ``inputs`` (None for
        failures) and ``errors`` sorted by index.

    Raises:
        ValueError: If ``memory=`` is passed via ``flow_kwargs``, or if
            ``concurrency < 1``.
        Exception: Re-raised when ``on_error="raise"`` and any input
            fails. The exception is the first one raised by a worker;
            other workers may already be cancelled when this surfaces.
    """
    if "memory" in flow_kwargs:
        raise ValueError(
            "batch_run does not support `memory=` in MVP. For "
            "memory-accumulating workflows write a serial loop instead: "
            "`for inp in inputs: await flow_fn(inp, cfg=cfg, memory=memory)`. "
            "See design doc §4.3.5."
        )
    if concurrency < 1:
        raise ValueError(f"concurrency must be >= 1, got {concurrency}")

    sem = asyncio.Semaphore(concurrency)
    results: list[OutputT | None] = [None] * len(inputs)
    errors: list[tuple[int, Exception]] = []
    started = time.monotonic()
    done_counter = 0

    async def run_one(i: int, inp: InputT) -> None: