向量化和重排在生产里很少是「一条文本、一次 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 决定要不要有限并发多个批次。顺序反了,限流和尾延迟会一起变差。
评论