Recipe: a hard dollar limit per flow run for LLM tasks (flow_run.id + retry_condition_fn)

[domondi1]

A flow that maps an LLM task over a batch (summarize 500 documents, enrich 2,000 rows, grade an eval set) can make a lot of paid calls, and task retries multiply them. Prefect bounds time and concurrency well, but there's no dollar ceiling for one flow run. This recipe adds one with an OpenAI client and prefect.runtime.flow_run.id: every task in the run sends the run id and a budget to a local gateway, which refuses a call with HTTP 402 before it reaches OpenAI once the run's budget is used.

Disclosure: I maintain the gateway (Inferrail, open source).

pip install prefect openai inferrail
export OPENAI_API_KEY=sk-...
inferrail serve --quickstart --app-mode      # http://127.0.0.1:8000

from openai import APIStatusError, OpenAI
from prefect import flow, task
from prefect.runtime import flow_run

def client_for_this_run(budget_usd: str) -> OpenAI:
    return OpenAI(
        base_url="http://127.0.0.1:8000/v1",
        api_key="unused",                     # the real key stays in the gateway
        default_headers={
            "X-Inferrail-Attribute-Work-Id": f"prefect-{flow_run.id}",   # one id per flow run
            "X-Inferrail-Budget-Usd": budget_usd,                        # its dollar ceiling
        },
    )

def retry_unless_budget_refusal(task, task_run, state) -> bool:
    try:
        state.result()
    except APIStatusError as e:
        return e.status_code != 402     # a budget refusal won't pass on retry
    except Exception:
        return True
    return True

@task(retries=2, retry_condition_fn=retry_unless_budget_refusal)
def summarize(doc: str, budget_usd: str) -> str:
    r = client_for_this_run(budget_usd).chat.completions.create(
        model="gpt-4o-mini",
        max_tokens=300,
        messages=[{"role": "user", "content": f"Summarize: {doc}"}],
    )
    return r.choices[0].message.content

@flow
def summarize_batch(docs: list[str], budget_usd: str = "5.00") -> list[str]:
    results = []
    for future in summarize.map(docs, budget_usd=budget_usd):
        try:
            results.append(future.result())
        except APIStatusError as e:
            results.append(f"skipped: budget ({e.status_code})")
    return results

After a run, inferrail work prefect-<flow run id> prints calls, tokens and dollars for it, and inferrail report --by work_id lists every run.

What I checked (prefect 3.8.7, openai 3.24.0, inferrail 0.4.15, local ephemeral Prefect, 8 mapped tasks against a stand-in OpenAI endpoint with fixed token counts):

• Enough budget: 8 tasks, 8 calls, all under prefect-<flow run id>.
• Budget for about half: 4 tasks ran, 4 were refused before reaching the upstream, including tasks running concurrently (reservations, not after-the-fact accounting), and the flow finished with the refused items marked.
• Retries: with a plain @task(retries=2), each refused task was retried twice, so 12 refused requests for 4 tasks. Refusals cost nothing, but it's wasted work. With retry_condition_fn as above: exactly 8 requests, and real transient errors still retry.

Caveats: a dollar budget needs a price for the model (most OpenAI and Anthropic models are built in, inferrail models; an unpriced one is refused, not guessed). The gateway stores usage and cost per call, never prompts or outputs. With workers on several machines, run one shared gateway and point base_url at it so the run's budget is shared across them.

If this is useful as a community recipe, I'm happy to host it as a standalone example.