用 Replicate API 构建可恢复的批量预测运行器：并发、限流与断点续跑

一次要生成几十张图时，需要控制并发、处理限流、中途失败后能接着跑，而且不能重复花钱。 这篇做 batch.py：读 CSV 里的提示词，最多同时跑 N 个预测，遇到 429 按返回的等待时间重试，每个预测用 Cancel-After 设服务端期限，结果写进 ledger.jsonl，重跑时跳过已成功的，并继续查询上次没结束的预测，而不是重新创建。

先看清成品

• 用无效 Token 真实运行：第一个请求返回 401 后立即停止，提示已完成的记录不会丢。
• 设计检查时发现的问题：第一版只在预测结束时写 ledger，中途中断后重跑会为“还在跑的”预测再建一次，重复计费。已改为创建后立刻写入，重跑时恢复查询。
• 用 mock 跑通 6 个场景（429、断点恢复、超时取消、aborted、401），pytest 6 项通过。

合规提醒：只用 Replicate 官方 HTTP API 和你自己的 API Token（replicate.com/account/api-tokens），Token 只放服务端，不共享账号。按运行时间计费，价格以 replicate.com/pricing 为准，本文不写价格。每个模型有自己的许可证和使用限制（在模型页面查看），生成内容要遵守 Replicate 服务条款，不生成侵权或违法内容，不冒用真人肖像。

运行说明：我没有 Replicate Token，没有跑过真实预测（也就没有产生费用）。真实跑过的是：用无效 Token 调用官方 API（2026-10-21 21:09–21:40，UTC+8，返回原样贴在文中），以及全部本地代码。正常的预测返回来自我写的 mock，字段照官方 HTTP 参考里的 prediction 对象。

做完之后你会得到什么？

需要准备哪些环境？

关键事实

文件结构

85-replicate-C-src/
├── batch.py
├── prompts.csv
└── test_batch.py

第 1 步：提示词

key,prompt
cat,一只橘猫坐在窗台上看雨，水彩画风格
tea,一杯冒着热气的绿茶，木桌，清晨的自然光
bike,停在胡同口的旧自行车，胶片质感
lamp,书桌上的台灯和摊开的笔记本，夜晚

第 2 步：批量脚本

"""批量跑 Replicate 预测：断点续跑（ledger.jsonl）、遇到 429 按提示等待、超时自动取消。
只用官方 HTTP API：创建预测、查询预测、取消预测。"""
import argparse, csv, json, os, re, sys, time
import requests

API = "https://api.replicate.com/v1"
UA = "binbho-tutorial/1.0"
TERMINAL = {"succeeded", "failed", "canceled", "aborted"}   # aborted 见“预测生命周期”页


def retry_after(resp, default=30):
    """429 的 body 形如 {"detail":"Request was throttled. Your rate limit resets in ~30s."}"""
    h = resp.headers.get("Retry-After")
    if h and h.isdigit():
        return int(h)
    m = re.search(r"resets in ~?(\d+)s", (resp.json() or {}).get("detail", ""))
    return int(m.group(1)) if m else default


class Runner:
    def __init__(self, token, model, ledger="ledger.jsonl", session=None, sleep=time.sleep, clock=time.monotonic):
        if not token.strip():
            raise SystemExit("缺少 REPLICATE_API_TOKEN")
        self.s = session or requests.Session()
        self.s.headers.update({"Authorization": f"Bearer {token.strip()}", "User-Agent": UA})
        self.model, self.ledger, self.sleep, self.clock = model, ledger, sleep, clock
        self.done = {}
        if os.path.exists(ledger):
            for line in open(ledger, encoding="utf-8"):
                e = json.loads(line); self.done[e["key"]] = e

    def _req(self, method, url, **kw):
        while True:
            r = self.s.request(method, url, timeout=60, **kw)
            if r.status_code == 429:
                w = retry_after(r); print(f"  429，等待 {w}s", file=sys.stderr); self.sleep(w); continue
            if r.status_code >= 400:
                d = r.json()
                raise RuntimeError(f"HTTP {r.status_code} {d.get('title', '')}: {d.get('detail')}")
            return r.json()

    def record(self, key, pred):
        e = {"key": key, "id": pred["id"], "status": pred["status"], "output": pred.get("output"),
             "error": pred.get("error"), "urls": pred.get("urls", {})}
        with open(self.ledger, "a", encoding="utf-8") as f:
            f.write(json.dumps(e, ensure_ascii=False) + "\n")
        self.done[key] = e

    def run(self, rows, max_inflight=2, timeout=300, interval=3):
        inflight = {}   # key -> (pred, started)
        # 上次中断时还没结束的预测：继续查询原来的 id，而不是重新创建（重新创建会再计费一次）
        for key, e in self.done.items():
            if e["status"] not in TERMINAL:
                inflight[key] = ({"id": e["id"], "status": e["status"], "urls": e["urls"]}, self.clock())
                print(f"  恢复 {key} → {e['id']}（{e['status']}）")
        todo = [r for r in rows if r["key"] not in inflight
                and self.done.get(r["key"], {}).get("status") != "succeeded"]
        print(f"共 {len(rows)} 条，已完成 {len(rows) - len(todo) - len(inflight)} 条，恢复 {len(inflight)} 条，新建 {len(todo)} 条")
        while todo or inflight:
            while todo and len(inflight) < max_inflight:
                r = todo.pop(0)
                # Cancel-After：服务端期限，客户端断开也会生效；没开始就超时 → aborted（不计费）
                p = self._req("POST", f"{API}/models/{self.model}/predictions", json={"input": {"prompt": r["prompt"]}},
                              headers={"Cancel-After": f"{timeout}s"})
                self.record(r["key"], p)          # 创建后立即落盘
                inflight[r["key"]] = (p, self.clock()); print(f"  创建 {r['key']} → {p['id']}")
            self.sleep(interval)
            for key, (p, t0) in list(inflight.items()):
                p = self._req("GET", p["urls"]["get"])
                if p["status"] in TERMINAL:
                    self.record(key, p); del inflight[key]; print(f"  {key}: {p['status']}")
                elif self.clock() - t0 > timeout:
                    p = self._req("POST", p["urls"]["cancel"])
                    self.record(key, {**p, "status": "canceled", "error": f"超过 {timeout}s，已取消"})
                    del inflight[key]; print(f"  {key}: 超时取消")
                else:
                    inflight[key] = (p, t0)
        bad = [k for k, e in self.done.items() if e["status"] != "succeeded"]
        print(f"完成。成功 {len(self.done) - len(bad)}，未成功 {bad or '无'}。输出链接 1 小时后失效，请及时下载（见 85A）。")
        return 0 if not bad else 2


def main(argv=None):
    sys.stdout.reconfigure(line_buffering=True)   # 和 stderr 的输出按时间顺序显示
    ap = argparse.ArgumentParser()
    ap.add_argument("csv"); ap.add_argument("--model", default="black-forest-labs/flux-schnell")
    ap.add_argument("--max-inflight", type=int, default=2); ap.add_argument("--timeout", type=int, default=300)
    a = ap.parse_args(argv)
    rows = list(csv.DictReader(open(a.csv, encoding="utf-8")))
    try:
        return Runner(os.environ.get("REPLICATE_API_TOKEN", ""), a.model).run(rows, a.max_inflight, a.timeout)
    except RuntimeError as e:
        print(f"中止：{e}（已完成的记录在 ledger.jsonl，修复后重跑会跳过它们）", file=sys.stderr); return 1


if __name__ == "__main__":
    sys.exit(main())

第 3 步：测试（mock）

"""假 Replicate（我写的 mock，prediction 字段按官方 HTTP 参考）。"""
import json
import pytest
import batch as b


class R:
    def __init__(self, code, body, headers=None):
        self.status_code, self._b, self.headers = code, body, headers or {}
    def json(self):
        return self._b


def pred(i, status, output=None):
    return {"id": i, "status": status, "output": output, "error": None,
            "urls": {"get": f"https://api.replicate.com/v1/predictions/{i}", "cancel": f"https://api.replicate.com/v1/predictions/{i}/cancel", "web": f"https://replicate.com/p/{i}"}}


class FakeReplicate:
    """每个预测第 2 次查询时成功；可以注入 429 和永远不结束的预测。"""
    def __init__(self, throttle_first=False, stuck=()):
        self.n, self.polls, self.calls, self.headers = 0, {}, [], {}
        self.throttle, self.stuck = throttle_first, set(stuck)
    def request(self, method, url, timeout=None, json=None, headers=None):
        self.calls.append((method, url))
        if method == "POST" and url.endswith("/predictions"):
            assert headers == {"Cancel-After": "300s"} or headers["Cancel-After"].endswith("s")
        if self.throttle:
            self.throttle = False
            return R(429, {"detail": "Request was throttled. Your rate limit resets in ~7s."})
        if method == "POST" and url.endswith("/predictions"):
            self.n += 1; pid = f"p{self.n}"
            if json["input"]["prompt"] in self.stuck:
                self.stuck.add(pid)
            return R(201, pred(pid, "starting"))
        pid = url.split("/predictions/")[1].split("/")[0]
        if url.endswith("/cancel"):
            return R(200, pred(pid, "canceled"))
        self.polls[pid] = self.polls.get(pid, 0) + 1
        if pid in self.stuck:
            return R(200, pred(pid, "processing"))
        return R(200, pred(pid, "succeeded" if self.polls[pid] >= 2 else "processing", [f"https://replicate.delivery/{pid}.webp"]))


ROWS = [{"key": k, "prompt": k} for k in ("a", "b", "c")]


def runner(tmp_path, fake, clock=None):
    t = [0]
    def tick():
        t[0] += 10; return t[0]
    return b.Runner("r8_x", "owner/model", ledger=str(tmp_path / "l.jsonl"), session=fake, sleep=lambda s: None, clock=clock or tick)


def test_all_succeed_and_429(tmp_path):
    f = FakeReplicate(throttle_first=True)
    assert runner(tmp_path, f).run(ROWS, max_inflight=2) == 0
    assert f.n == 3


def test_retry_after_parsing():
    assert b.retry_after(R(429, {"detail": "Request was throttled. Your rate limit resets in ~30s."})) == 30
    assert b.retry_after(R(429, {}, {"Retry-After": "5"})) == 5


def test_resume_skips_done_and_does_not_recreate_inflight(tmp_path):
    led = tmp_path / "l.jsonl"
    led.write_text("\n".join(json.dumps(e) for e in [
        {"key": "a", "id": "old1", "status": "succeeded", "output": [], "error": None, "urls": pred("old1", "")["urls"]},
        {"key": "b", "id": "old2", "status": "processing", "output": None, "error": None, "urls": pred("old2", "")["urls"]}]) + "\n")
    f = FakeReplicate()
    assert runner(tmp_path, f).run(ROWS) == 0
    assert f.n == 1                                  # 只为 c 新建
    assert ("GET", "https://api.replicate.com/v1/predictions/old2") in f.calls


def test_timeout_cancels(tmp_path):
    f = FakeReplicate(stuck={"b"})
    assert runner(tmp_path, f).run(ROWS, timeout=25) == 2
    assert any(u.endswith("/cancel") for m, u in f.calls)


def test_401_stops(tmp_path):
    class Bad(FakeReplicate):
        def request(self, *a, **k):
            return R(401, {"status": 401, "detail": "Invalid or expired API token.", "title": "Authentication required"})
    with pytest.raises(RuntimeError, match="401"):
        runner(tmp_path, Bad()).run(ROWS)


def test_aborted_is_terminal(tmp_path):
    class Ab(FakeReplicate):
        def request(self, method, url, timeout=None, json=None, headers=None):
            if method == "GET":
                return R(200, pred(url.rsplit("/", 1)[1], "aborted"))
            return super().request(method, url, timeout, json, headers)
    assert runner(tmp_path, Ab()).run(ROWS[:1]) == 2

第 4 步：运行

$ REPLICATE_API_TOKEN=r8_invalid_demo python batch.py prompts.csv   # 2026-10-21 21:40 UTC+8
共 4 条，已完成 0 条，恢复 0 条，新建 4 条
中止：HTTP 401 Authentication required: Invalid or expired API token.（已完成的记录在 ledger.jsonl，修复后重跑会跳过它们）
$ echo $?
1
$ python -m pytest -q
6 passed

# 有 Token 时的输出格式（我没跑过）：
共 4 条，已完成 0 条，恢复 0 条，新建 4 条
  创建 cat → <id>
  创建 tea → <id>
  cat: succeeded
  …
完成。成功 4，未成功 无。输出链接 1 小时后失效，请及时下载（见 85A）。

怎么检查结果？

• 先用 2–3 条提示词试跑，确认单张耗时和费用，再放大。
• 中途按 Ctrl+C 停掉，再运行一次：应该看到“恢复 … 条”，而不是重新创建。
• ledger.jsonl 里同一个 key 可能有多行（创建时一行、结束时一行），以最后一行为准。
• 有赠送额度但没绑卡时，官方限制为每分钟最多 6 次，--max-inflight 设成 1。

调试示例

1. 中断后重跑会重复计费。 第一版的 ledger 只在预测进入终态时写入。我在写“断点续跑”的测试时想到这个场景：创建了 2 个预测，第 3 个请求时网络中断或 401，脚本退出。此时那 2 个预测还在 Replicate 上运行（也在计费），但 ledger 里没有它们；重跑时会被当成“没做过”，再创建一次。修复：

• 创建成功后立刻 record()，状态是 starting；
• 启动时把 ledger 里非终态的预测放回 inflight，用原来的 urls.get 继续查询；
• 测试 test_resume_skips_done_and_does_not_recreate_inflight 检查：恢复的 old2 只被查询、没被重建，只有 c 是新建的。

2. 客户端超时取消不够可靠。 原来只在客户端计时、超时后调用 cancel；如果脚本自己挂了，预测会一直跑到平台上限（官方说 30 分钟）。改为创建时带 Cancel-After: {timeout}s，让服务端到期自动停止；没开始就到期的会是 aborted 而不收费。客户端的 cancel 保留作为兜底，aborted 也加入了终态（见 A 篇调试示例）。

3. 输出顺序错乱。 第一次用无效 Token 运行时，stderr 的“中止：…”出现在 stdout 的“共 4 条…”之前，因为输出被重定向时 stdout 有缓冲。在 main() 里设置行缓冲后顺序正常。

常见问题

模板规则

• 创建后立刻落盘。
• 重跑时恢复查询，不重复创建。
• 设服务端期限 Cancel-After。
• 先小批量试跑看费用。

项目资源

批量任务最怕的不是失败，而是失败之后重复花钱。