Batch steps

Use BatchStep when a step does the same expensive call for many items. Each item's result is cached on its own, so a rerun after a failure computes only the missing items.

Implement three methods:

1. items(**inputs): return one BatchItem(key=..., data=...) per item.
2. process_item(item, **inputs): compute one item.
3. assemble(results, **inputs): combine the per-item results into the step result.

Instead of items(), you can implement items_data() and return plain data. adag then derives each item key from three things: the step's own hash (class, code, fields), the item data, and every other input of the step.

Set items_source_keys to the name of the input that holds the items, for example items_source_keys = ("names",). Without it, the whole list is part of every item key, so adding one item recomputes all of them. The example in Cache stores sets it.

Override process_batch() to process items concurrently.

"""BatchStep: each item is cached on its own, so a rerun resumes after a failure."""

import tempfile
from dataclasses import dataclass
from pathlib import Path

from pydantic import BaseModel

from adag import BatchItem, BatchStep, InputParam, NullObserver, Pipeline

CALLS: list[str] = []
FAIL_ON: set[str] = {"carol"}


class Greeting(BaseModel):
    text: str


class Greetings(BaseModel):
    lines: list[str]


@dataclass(frozen=True)
class Greet(BatchStep[Greetings, Greeting]):
    def items(self, names: list[str]) -> list[BatchItem]:
        return [BatchItem(key=name, data=name) for name in names]

    def process_item(self, item: BatchItem, names: list[str]) -> Greeting:
        CALLS.append(item.data)
        if item.data in FAIL_ON:
            raise RuntimeError(f"API error for {item.data}")
        return Greeting(text=f"hello {item.data}")

    def assemble(self, results: list[Greeting], names: list[str]) -> Greetings:
        return Greetings(lines=[result.text for result in results])


names = InputParam[list[str]]("names")
pipeline = Pipeline(Greet()(names=names), cache_dir=Path(tempfile.mkdtemp()))
everyone = ["alice", "bob", "carol", "dave"]

try:
    pipeline.execute(observer=NullObserver(), names=everyone)
except RuntimeError as error:
    print("first run failed:", error, "| called:", CALLS)

CALLS.clear()
FAIL_ON.clear()
result = pipeline.execute(observer=NullObserver(), names=everyone)
print("second run called:", CALLS)
print(result.lines)

first run failed: API error for carol | called: ['alice', 'bob', 'carol']
second run called: ['carol', 'dave']
['hello alice', 'hello bob', 'hello carol', 'hello dave']

flowchart TB
    items["items(): one BatchItem per item"] --> each{"Item result saved?"}
    each -- yes --> load["Load it"]
    each -- no --> proc["process_item(): compute it"] --> save["Save it right away"]
    load --> asm["assemble(): combine all results"]
    save --> asm
    asm --> done(["Step result"])

A crash loses only the item in progress: items already saved are loaded on the next run.

Item results live under <cache_dir>/.items/<namespace>/. The namespace is the class name in snake case (greet above).