向量化和重排在生产里很少是「一条文本、一次 HTTP」。把短请求在时延预算内合并成有界微批,能同时压住供应商 QPS、连接数和尾延迟。本文给出按模型切分的 Node.js 实现:超时强制刷新、失败拆批重试,并接上取消信号、幂等键和基础指标。

逐条调用会同时打满 QPS 和尾延迟

在线检索链路里,embedding 和 rerank 的入参几乎总是短文本:改写后的 query、召回的 chunk、多路候选再打一遍分数。本地调试用 for 循环逐条 fetch 最清楚,但这个形状一旦进生产,成本和延迟会叠在同一处。

供应商计的是请求次数和 token,不是「你有多短」。一千条平均只有几十 token 的句子,逐条调用就是一千次 HTTP、一千次鉴权、一千次进入限流桶;改成一次请求里的 input / documents 数组,同类短文本往往能被同一趟带走。TLS 握手、连接占用、事件循环上悬挂的 Promise 数量,同样跟请求次数走,不跟单条文本长度走。

尾延迟更不容易从平均值里看出来。逐条调用的 P99 接近「最慢那一次 + 排队」。检索接口通常还有自己的总时延预算,embedding 把预算吃光后,rerank 只能在剩余时间里挨打。把并发用信号量拉到几十,只会更快碰到 429,然后在重试抖动里把尾延迟再推高一截。

所以问题不是「要不要并发」,而是请求形状错了:本可以在几毫秒窗口里合并的同类短文本,被拆成了供应商最不划算、你这边连接压力最大的调用方式。微批要做的是把形状改回来,而不是再加一层盲目并发。

有界微批:按模型切分,在时延预算内合并

能合并的前提是一次 HTTP 能表达这批数据。桶键必须落在供应商的请求边界上:embedding 按模型切开,text-embedding-3-small 和 text-embedding-3-large 不能混发;rerank 还要加上 query,不同用户的查询塞进同一 documents 没有语义,分数也对不回原请求。

「有界」指三条硬限制,而不是尽量多攒一点。

条数上限跟官方请求限制对齐。OpenAI Embeddings 的 input 可以是字符串数组,单次最多 2048 条,同时还受总 token 约束。Rerank 的 documents 同样有上限,实现里应取文档中的上限再留余量,例如先按 64 或 128 截断,而不是用线上拍出来的「经验吞吐」。时延上限是检索预算里能让出的那一小段,常见是 8 到 20 毫秒:窗口过大,首条请求会明显变慢;过小,批次经常是 1,微批等于没做。具体取值要用你们自己的链路预算去量,而不是抄一个「经验数字」当结论。

第三条是取消。上游 AbortSignal 一旦触发,这条不得再进入即将发出的批次;已经发出的 fetch 要能被信号打断,且不能把同批其他人 reject 掉。刷新条件只保留三种:桶满、等待到期、关闭前 flush。再加「空闲探测」会让延迟分布无法解释,排障时你也不知道一次调用为什么提前走了。

Node.js 实现:队列、刷新、取消与拆批

下面这段 ESM 可直接保存为 micro-batch.mjs,Node.js 18 及以上运行。核心类不绑供应商,只负责分桶、超时、abort,以及 runBatch 失败时二分拆开重试。无密钥时执行 node micro-batch.mjs 走本地模拟,用来看合并和拆批;node micro-batch.mjs openai 且配置 OPENAI_API_KEY 则打真实 Embeddings 接口。

// micro-batch.mjs
import { randomUUID } from "node:crypto";

export const metrics = {
  flushSize: 0,
  flushTimeout: 0,
  flushShutdown: 0,
  splitRetries: 0,
  cancelled: 0,
  batches: 0,
  items: 0,
};

export class BoundedMicroBatcher {
  constructor({ maxBatchSize, maxWaitMs, keyFn, runBatch }) {
    if (maxBatchSize < 1) throw new Error("maxBatchSize must be >= 1");
    if (maxWaitMs < 0) throw new Error("maxWaitMs must be >= 0");
    this.maxBatchSize = maxBatchSize;
    this.maxWaitMs = maxWaitMs;
    this.keyFn = keyFn;
    this.runBatch = runBatch;
    this.buckets = new Map();
    this.closed = false;
  }

  submit(item, { signal } = {}) {
    if (this.closed) return Promise.reject(new Error("batcher closed"));
    if (signal?.aborted) {
      metrics.cancelled += 1;
      return Promise.reject(signal.reason ?? new Error("aborted"));
    }

    return new Promise((resolve, reject) => {
      const key = this.keyFn(item);
      const waiter = { item, resolve, reject, signal, onAbort: null };

      const abortNow = () => {
        const bucket = this.buckets.get(key);
        if (bucket) {
          const idx = bucket.waiters.indexOf(waiter);
          if (idx >= 0) {
            bucket.waiters.splice(idx, 1);
            metrics.cancelled += 1;
            if (bucket.waiters.length === 0) {
              if (bucket.timer) clearTimeout(bucket.timer);
              this.buckets.delete(key);
            }
          }
        }
        reject(signal.reason ?? new Error("aborted"));
      };

      if (signal) {
        waiter.onAbort = abortNow;
        signal.addEventListener("abort", abortNow, { once: true });
      }

      let bucket = this.buckets.get(key);
      if (!bucket) {
        bucket = { waiters: [], timer: null, firstAt: Date.now() };
        this.buckets.set(key, bucket);
        if (this.maxWaitMs > 0) {
          bucket.timer = setTimeout(() => {
            bucket.timer = null;
            void this.flushKey(key, "timeout");
          }, this.maxWaitMs);
        }
      }
      bucket.waiters.push(waiter);
      if (bucket.waiters.length >= this.maxBatchSize) {
        void this.flushKey(key, "size");
      }
    });
  }

  async flush() {
    await Promise.all([...this.buckets.keys()].map((k) => this.flushKey(k, "shutdown")));
  }

  async close() {
    this.closed = true;
    await this.flush();
  }

  async flushKey(key, reason) {
    const bucket = this.buckets.get(key);
    if (!bucket || bucket.waiters.length === 0) return;
    this.buckets.delete(key);
    if (bucket.timer) {
      clearTimeout(bucket.timer);
      bucket.timer = null;
    }
    if (reason === "size") metrics.flushSize += 1;
    else if (reason === "timeout") metrics.flushTimeout += 1;
    else metrics.flushShutdown += 1;

    const waiters = bucket.waiters;
    for (const w of waiters) {
      if (w.signal && w.onAbort) w.signal.removeEventListener("abort", w.onAbort);
    }
    metrics.batches += 1;
    metrics.items += waiters.length;
    await this.runWithSplit(waiters, key, reason, Date.now() - bucket.firstAt);
  }

  async runWithSplit(waiters, key, reason, waitMs) {
    const alive = waiters.filter((w) => !w.signal?.aborted);
    for (const w of waiters) {
      if (w.signal?.aborted) {
        metrics.cancelled += 1;
        w.reject(w.signal.reason ?? new Error("aborted"));
      }
    }
    if (alive.length === 0) return;

    const idempotencyKey = randomUUID();
    try {
      const results = await this.runBatch(
        alive.map((w) => w.item),
        { signal: undefined, idempotencyKey, key, reason, waitMs }
      );
      if (!Array.isArray(results) || results.length !== alive.length) {
        throw new Error("runBatch must return an array aligned with items");
      }
      alive.forEach((w, i) => w.resolve(results[i]));
    } catch (err) {
      if (alive.length === 1) {
        alive[0].reject(err);
        return;
      }
      metrics.splitRetries += 1;
      const mid = Math.ceil(alive.length / 2);
      await this.runWithSplit(alive.slice(0, mid), key, reason, waitMs);
      await this.runWithSplit(alive.slice(mid), key, reason, waitMs);
    }
  }
}

function isRetryableStatus(status) {
  return status === 429 || status >= 500;
}

export async function openaiEmbedBatch(texts, { model, signal, idempotencyKey }) {
  const apiKey = process.env.OPENAI_API_KEY;
  if (!apiKey) throw new Error("OPENAI_API_KEY required");
  const res = await fetch("https://api.openai.com/v1/embeddings", {
    method: "POST",
    signal,
    headers: {
      Authorization: `Bearer ${apiKey}`,
      "Content-Type": "application/json",
      "Idempotency-Key": idempotencyKey,
    },
    body: JSON.stringify({ model, input: texts }),
  });
  if (!res.ok) {
    const body = await res.text();
    const err = new Error(`openai embeddings ${res.status}: ${body.slice(0, 200)}`);
    err.status = res.status;
    err.retryable = isRetryableStatus(res.status);
    throw err;
  }
  const json = await res.json();
  const byIndex = new Map(json.data.map((row) => [row.index, row.embedding]));
  return texts.map((_, i) => byIndex.get(i));
}

export async function cohereRerankBatch(query, documents, { model, topN, signal }) {
  const apiKey = process.env.COHERE_API_KEY;
  if (!apiKey) throw new Error("COHERE_API_KEY required");
  const res = await fetch("https://api.cohere.com/v1/rerank", {
    method: "POST",
    signal,
    headers: {
      Authorization: `Bearer ${apiKey}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify({
      model,
      query,
      documents,
      top_n: topN ?? documents.length,
      return_documents: false,
    }),
  });
  if (!res.ok) {
    const body = await res.text();
    const err = new Error(`cohere rerank ${res.status}: ${body.slice(0, 200)}`);
    err.status = res.status;
    err.retryable = isRetryableStatus(res.status);
    throw err;
  }
  const json = await res.json();
  return json.results;
}

async function demoLocal() {
  const batcher = new BoundedMicroBatcher({
    maxBatchSize: 4,
    maxWaitMs: 15,
    keyFn: (item) => item.model,
    runBatch: async (items, ctx) => {
      console.log("flush", {
        key: ctx.key,
        reason: ctx.reason,
        size: items.length,
        waitMs: ctx.waitMs,
        idempotencyKey: ctx.idempotencyKey,
      });
      if (items.some((it) => it.text === "boom")) {
        throw new Error("provider rejected batch");
      }
      return items.map((it) => ({ dimHint: it.text.length, model: it.model }));
    },
  });

  const merged = Promise.all(
    ["alpha", "beta", "gamma", "delta", "epsilon"].map((text) =>
      batcher.submit({ model: "text-embedding-3-small", text })
    )
  );
  const otherModel = batcher.submit({
    model: "text-embedding-3-large",
    text: "only-large",
  });
  const ac = new AbortController();
  const cancelled = batcher
    .submit({ model: "text-embedding-3-small", text: "never" }, { signal: ac.signal })
    .catch((err) => ({ cancelled: true, message: String(err.message) }));
  ac.abort();
  const split = await Promise.all(
    ["ok-1", "boom", "ok-2"].map((text) =>
      batcher
        .submit({ model: "split-demo", text })
        .then((value) => ({ text, ok: true, value }))
        .catch((err) => ({ text, ok: false, error: err.message }))
    )
  );

  console.log({
    merged: await merged,
    otherModel: await otherModel,
    cancelled: await cancelled,
    split,
    metrics,
  });
  await batcher.close();
}

async function demoOpenAI() {
  const model = "text-embedding-3-small";
  const batcher = new BoundedMicroBatcher({
    maxBatchSize: 16,
    maxWaitMs: 20,
    keyFn: (item) => `embed:${item.model}`,
    runBatch: (items, ctx) =>
      openaiEmbedBatch(
        items.map((it) => it.text),
        {
          model: items[0].model,
          signal: ctx.signal,
          idempotencyKey: ctx.idempotencyKey,
        }
      ),
  });
  const texts = ["微批合并", "向量检索", "重排模型"];
  const vectors = await Promise.all(texts.map((text) => batcher.submit({ model, text })));
  await batcher.close();
  console.log(vectors.map((vec, i) => ({ text: texts[i], dim: vec.length })));
}

const mode = process.argv[2] ?? "local";
if (mode === "openai") {
  await demoOpenAI();
} else {
  await demoLocal();
}

submit 按 keyFn 入桶。桶内第一条启动定时器,人满或到期则 flush。abort 发生在发出前,会从桶中摘除自己,不影响他人。runBatch 必须返回和入参等长的数组,调用方才能按条 resolve。

拆批是为了「一批里有一条非法输入」这类局部失败:二分直到单条,成功的那一半仍会返回。429 和 5xx 不该走拆批,否则一次限流会变成多次限流。示例把策略留给 runBatch 的抛错方式;生产应在适配器里区分 status,限流则整批退避并尊重 Retry-After。

接到 Embedding 与 Rerank

Embeddings 的合并点就是 input 数组,分桶键用 embed:${model}。OpenAI 返回的 data[].index 要按入参顺序重排,不要假设数组顺序稳定。幂等键按「这一次批次尝试」生成:同一尝试的网络重放沿用同一把键;拆批后的子批必须换新键,避免服务端把子批当成原批的重复。OpenAI 的 POST 接受 Idempotency-Key;它不能替代你在客户端做去重,但能避免同一次尝试被网关重放成两次计费。

Rerank 不要跨 query 合并。分桶键用 rerank:${model}:${query},每条 submit 只带一个 document,flush 时组装 documents 调一次 Cohere v1/rerank。top_n 建议等于本批条数,这样才能给每个 document 拿回 relevance_score,再按原 index 写回。不同 query 会自然落在不同桶,各自受 maxWaitMs 约束,不会为了凑批而串台。接入时不要另起一套队列,直接复用上面的类:

export function createRerankBatcher({ maxBatchSize = 32, maxWaitMs = 15 } = {}) {
  return new BoundedMicroBatcher({
    maxBatchSize,
    maxWaitMs,
    keyFn: (item) => `rerank:${item.model}:${item.query}`,
    runBatch: async (items, ctx) => {
      const documents = items.map((it) => it.document);
      const results = await cohereRerankBatch(items[0].query, documents, {
        model: items[0].model,
        topN: documents.length,
        signal: ctx.signal,
      });
      const scoreByIndex = new Map(results.map((row) => [row.index, row.relevance_score]));
      return documents.map((_, i) => scoreByIndex.get(i) ?? 0);
    },
  });
}

观测至少要能回答三个问题:窗口是否过大(timeout 刷新为主且 batch_size 接近 1)、是否过小(size 触发频繁但上游仍 429)、拆批是否在帮倒忙(splitRetries 和 429 一起涨)。flush_reason、batch_size、从首条入队到发出的 wait_ms、cancelled 这几项,比一条笼统的 duration 有用。指标先挂在进程内计数器上就够排查;再往 Prometheus 里送时,按 model 和 reason 打标签即可,不必一上来上分布式追踪。

总结

生产里 embedding 和 rerank 的瓶颈,经常是请求被拆得太碎,而不是模型本身算得慢。按模型以及 rerank 的 query 分桶,用条数和等待双上限强制刷新,用 AbortSignal 贯穿入队到发出前的摘除,失败时按错误类型选择整批退避或拆批,再用幂等键和分原因指标把行为固定下来。先让同一模型的短文本在时延预算内稳定合并,再视 RPM 决定要不要有限并发多个批次。顺序反了,限流和尾延迟会一起变差。