生产级 AI 推理网关:高并发请求调度、请求去重与异步取消的工程实践

生产级 AI 推理网关:高并发请求调度、请求去重与异步取消的工程实践

在大规模 AI Agent 部署中,推理引擎(vLLM/SGLang/TGI)解决的是"怎么算得快",而推理网关解决的是"怎么管得住"。当一个生产系统要在 1000 QPS 下稳定运行 200+ 种不同模型的推理请求时,网关层的设计决定了整个系统的稳定性边界。


一、为什么需要推理网关:从裸 Engine 到生产化

1.1 裸推理引擎的局限性

现代推理引擎(vLLM、SGLang、TensorRT-LLM 等)在单请求/低并发场景下性能优异,但直接暴露在业务流量中时面临几个核心问题:

请求无序性:HTTP/gRPC 请求直接打到 Engine,无法区分优先级。一个 batch 的超时请求会阻塞后续高优先级请求。

资源碎片化:Engine 的 Continuous Batching 机制对混合长度请求敏感。长短请求混合时,长尾请求占用的 KV Cache 阻塞新请求入队。

缺乏全局限流:单机 Engine 感知不到集群全局的负载。A 实例空闲而 B 实例过载时,请求仍在被路由到 B。

1.2 网关的核心职责

一个生产级推理网关(如 SGLang Router、 Dynamo Router、llm-d、Envoy AI Gateway)承担以下职责:

职责 描述 不做的后果
智能路由 根据负载/模型/版本分发请求 单实例过载,延迟飙升
请求去重 相同 prompt 复用已有结果 重复计费 + 资源浪费
超时取消 优雅终止过期请求 Goroutine/Future 泄漏
优先级调度 按用户等级/任务类型排序 VIP 用户被低优请求阻塞
批处理聚合 小请求聚合成大 batch 吞吐下降

二、请求路由:从轮询到感知调度

2.1 传统负载均衡的局限性

经典的 Round-Robin、Least-Connections、一致性哈希在推理场景下都会出问题:

  • Round-Robin:无视实例当前 load,可能把请求路由到正在生成 4096 token 长回复的实例
  • Least-Connections:不考虑请求的计算量——1 个 32K 上下文请求的计算量远大于 10 个短请求
  • 一致性哈希:适用于缓存亲和,但不反映 KV Cache 使用率

2.2 基于可观测指标的感知调度(v1)

生产环境中最实用的方案是"基于可观测指标的加权路由":

/// 推理实例的可观测状态
#[derive(Debug, Clone, Default)]
struct InstanceMetrics {
    /// 当前排队请求数
    queue_depth: usize,
    /// 正在生成的批大小
    current_batch_size: usize,
    /// KV Cache 占用率 (0.0-1.0)
    kv_cache_utilization: f32,
    /// 平均首 token 延迟(毫秒)
    avg_ttft_ms: f64,
    /// 平均生成延迟(毫秒/token)
    avg_tpot_ms: f64,
    /// 当前 GPU 显存使用率
    gpu_memory_utilization: f32,
}

/// 综合评分(越低表示负载越轻)
fn instance_score(m: &InstanceMetrics) -> f64 {
    // 排队请求——最直接的负载指标
    let queue_score = m.queue_depth as f64 * 1.0;

    // KV Cache 压力——超过 80% 急剧恶化
    let kv_score = if m.kv_cache_utilization > 0.8 {
        m.kv_cache_utilization as f64 * 50.0  // 高压惩罚
    } else {
        m.kv_cache_utilization as f64 * 2.0
    };

    // 批处理利用率——接近上限时不加新请求
    let batch_score = if m.current_batch_size >= MAX_BATCH_SIZE {
        1000.0 // 饱和实例,不在路由
    } else {
        m.current_batch_size as f64 * 0.5
    };

    // 结合 TTFT 历史(EWMA 平滑)
    let latency_score = m.avg_ttft_ms / 100.0;

    queue_score + kv_score + batch_score + latency_score
}

/// 选择得分最低的实例
fn select_instance(instances: &[InstanceMetrics]) -> Option<usize> {
    instances.iter()
        .enumerate()
        .filter(|(_, m)| m.current_batch_size < MAX_BATCH_SIZE)
        .min_by(|(_, a), (_, b)| {
            instance_score(a).partial_cmp(&instance_score(b)).unwrap()
        })
        .map(|(id, _)| id)
}

2.3 SGLang Router 的 Prefix-Aware Routing

对于共享系统 Prompt 的请求(大多数 Production 场景),Prefix-Aware Routing 是更聪明的策略:

use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

/// 使用前缀哈希将"类似请求"路由到同一实例
/// 这样目标实例的 KV Cache 命中率高,减少重复计算
fn prefix_aware_route(
    prompt: &str,
    prefix_length: usize,  // 通常取 system prompt 长度
    instance_count: usize,
    instances: &[InstanceMetrics],
) -> Option<usize> {
    // 取 system prompt 的前缀
    let prefix = &prompt[..prompt.len().min(prefix_length)];

    let mut hasher = DefaultHasher::new();
    prefix.hash(&mut hasher);
    let hash = hasher.finish();

    // 使用一致性哈希确定"候选实例"
    let candidate = hash as usize % instance_count;

    // 检查候选实例是否过载
    if instances[candidate].kv_cache_utilization < 0.85 
        && instances[candidate].current_batch_size < MAX_BATCH_SIZE * 3 / 4 {
        Some(candidate)
    } else {
        // 失败时回退到一般评分路由
        select_instance(instances)
    }
}

这个设计在 Agent 场景下特别有效,因为多个 Agent 往往共享相同的 System Prompt(角色设定/工具定义),prefix-aware 路由使 KV Cache 命中率提升 30-60%。


三、请求去重:Semantic Cache 与精确去重

3.1 Deduplication vs Semantic Caching

请求去重有两层:

  1. 精确去重:完全相同的 prompt(包括 temperature/max_tokens 等参数)→ 直接返回缓存结果
  2. 语义去重:语义相同但字面不同("天气如何" vs "今天气温多少度")→ 通过 embedding 相似度匹配

精确去重在生产 Agent 场景下远比想象中常见——同一份文档的 RAG 查询可能来自数百个并发用户。

3.2 精确去重的工程实现

use tokio::sync::RwLock;
use lru::LruCache;
use std::num::NonZeroUsize;
use std::sync::Arc;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};

/// 请求指纹(用于精确去重)
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
struct RequestFingerprint {
    model_id: String,
    prompt_hash: String,
    temperature: u32,        // 放大 1000 倍存储 (0.7 → 700)
    max_tokens: u32,
    stop_sequences_hash: String,
}

impl RequestFingerprint {
    fn new(model_id: &str, prompt: &str, temperature: f32, max_tokens: u32, stop_seqs: &[String]) -> Self {
        let mut hasher = blake3::Hasher::new();
        hasher.update(prompt.as_bytes());
        let prompt_hash = URL_SAFE_NO_PAD.encode(hasher.finalize().as_bytes());

        let mut stop_hasher = blake3::Hasher::new();
        for s in stop_seqs {
            stop_hasher.update(s.as_bytes());
        }

        Self {
            model_id: model_id.to_string(),
            prompt_hash,
            temperature: (temperature * 1000.0) as u32,
            max_tokens,
            stop_sequences_hash: URL_SAFE_NO_PAD.encode(stop_hasher.finalize().as_bytes()),
        }
    }
}

/// 去重缓存:相同请求的响应可以被"携带"
struct DeduplicationCache {
    /// 等待中的请求(正在处理的相同请求)
    in_flight: Arc<RwLock<HashMap<RequestFingerprint, tokio::sync::watch::Receiver<Arc<String>>>>>,
    /// LRU 缓存已完成的结果(供后续相同请求直接使用)
    completed: Arc<RwLock<LruCache<RequestFingerprint, Arc<String>>>>,
}

impl DeduplicationCache {
    async fn get_or_compute<F, Fut>(
        &self,
        fp: RequestFingerprint,
        compute: F,
    ) -> Arc<String>
    where
        F: FnOnce() -> Fut,
        Fut: std::future::Future<Output = String> + Send + 'static,
    {
        // 1. 先检查 LRU 已完成缓存
        {
            let cache = self.completed.read().await;
            if let Some(result) = cache.peek(&fp) {
                metrics::counter!("dedup.lru_hit").increment(1);
                return result.clone();
            }
        }

        // 2. 再检查是否有相同请求 in-flight
        {
            let mut in_flight = self.in_flight.write().await;
            if let Some(rx) = in_flight.get(&fp) {
                let rx = rx.clone(); // watch::Receiver 是 Clone 的
                metrics::counter!("dedup.inflight_hit").increment(1);
                in_flight.remove(&fp); // 写完立即释放写锁
                drop(in_flight);
                // 等待结果
                let _ = rx.clone().changed().await;
                return rx.borrow().clone();
            }
        }

        // 3. 没有缓存,开始计算,注册 in-flight
        let (tx, rx) = tokio::sync::watch::channel(Arc::new(String::new()));
        {
            let mut in_flight = self.in_flight.write().await;
            in_flight.insert(fp.clone(), rx);
        }

        // 执行实际推理
        let result = compute().await;
        let result = Arc::new(result);

        // 通知所有等待者
        let _ = tx.send(result.clone());

        // 从 in-flight 移除
        {
            let mut in_flight = self.in_flight.write().await;
            in_flight.remove(&fp);
        }

        // 写入 LRU(给后续请求用)
        {
            let mut completed = self.completed.write().await;
            completed.put(fp, result.clone());
        }

        result
    }
}

3.3 为什么不用 redis/Memcached?

精确去重有三个硬约束让传统 KV Store 不太合适:

  1. 毫秒级过期:去重窗口通常只有 5-30 秒(请求处理期间),Redis 的 TTL 精度不够
  2. in-flight 等待:需要"请求合并"(Request Coalescing)——第二个相同请求需要等待第一个完成然后共享结果,这是 Pub/Sub 能做但自带复杂性
  3. 响应体很大:一个推理响应可能 1-50KB,反复序列化/网络传输的代价 > 直接传 Arc

最佳实践是进程内 LRU + Tokio watch channel 的 in-flight 合并。多实例间共享去重状态是另一个复杂话题,通常通过 prefix-aware routing(让相同请求落到同一实例)来避免。


四、超时与取消:Rust Async 的双刃剑

4.1 为什么推理取消特别难

推理请求的取消不是简单的"drop 一个 Future"就结束了。当一个流式推理请求被取消时,需要:

  1. 通知 Engine 停止生成(释放 KV Cache 块)
  2. 回收 GPU 显存
  3. 关闭 SSE/gRPC 流
  4. 释放连接/文件描述符
  5. 记录取消指标(用于计费/监控)

如果这些步骤中有任何遗漏,就会导致"幽灵续航"——客户端已经超时了,但 GPU 还在不知疲倦地生成用户永远不会看到的 token。

4.2 Rust 的 Async Cancellation Safety 问题

Rust 的 async/await 模型在简洁的代价下隐藏了 cancel safety 陷阱。考虑以下看起来正确的代码:

// ❌ 危险代码:tokio::select! 取消不安全
async fn unsafe_inference(handler: &mut InferenceHandler, request: Request) -> Result<String> {
    // 过程1:锁定引擎槽位
    let slot = handler.acquire_slot().await?;
    // 如果在这里被取消(外层 select),slot 永远不会被释放!

    // 过程2:提交推理请求
    let result = handler.run_inference(request).await?;

    // 过程3:释放槽位
    drop(slot);
    Ok(result)
}

问题:tokio::select! 在超时分支执行后,会 drop 未完成的 Future。如果 Future 在 acquire_slot 和 drop(slot) 之间被 cancel,slot 资源就会被泄漏。

4.3 正确的 Cancel Safety 模式

/// 模式1:用 RAII Guard 保护关键资源
struct InferenceSlotGuard {
    inner: Option<Box<InferenceSlot>>,  // Some=持有,None=已释放
}

impl InferenceSlotGuard {
    fn new(slot: Box<InferenceSlot>) -> Self {
        Self { inner: Some(slot) }
    }

    fn release(mut self) -> Box<InferenceSlot> {
        self.inner.take().expect("slot already released")
    }
}

impl Drop for InferenceSlotGuard {
    fn drop(&mut self) {
        if let Some(slot) = self.inner.take() {
            // 释放槽位到池
            SLOT_POOL.push(slot);
            metrics::counter!("slot.released").increment(1);
        }
    }
}

/// 将可能 cancel 的操作封装为非 async 的普通函数
async fn safe_inference(handler: &mut InferenceHandler, request: Request) -> Result<String> {
    // SAFE: acquire_slot 是原子的,不可取消
    let slot = handler.acquire_slot().await?;
    let _guard = InferenceSlotGuard::new(slot);  // RAII 保护

    // 推理:可被 cancel,但 slot 一定会通过 _guard 释放
    let result = handler.run_inference(request).await?;

    // 显式释放(选择不自动 Drop)
    let slot = _guard.release();
    drop(slot);
    Ok(result)
}

4.4 更完善的设计:分层取消

在实际生产中,不同层级的取消需要不同的处理策略:

/// 取消策略
#[derive(Debug, Clone, Copy)]
enum CancelPolicy {
    /// 立即释放所有资源,不回写缓存
    Immediate,
    /// 将已生成部分写入 cache,再释放
    CachePartial,
    /// 不取消,让推理跑完并缓存结果
    /// (在"这个请求很快会被重发"的预测下有用)
    FinishAndCache,
}

/// 推理任务上下文
struct InferenceTask {
    id: Uuid,
    cancel_policy: CancelPolicy,
    /// 用 tokio_util::sync::CancellationToken 作为取消信号源
    cancel_token: CancellationToken,
}

impl InferenceTask {
    async fn run(&self, request: Request, engine: &dyn InferenceEngine) -> Result<InferenceResult> {
        // 子上下文:超时控制
        let deadline = Instant::now() + Duration::from_secs(request.timeout_secs);

        tokio::select! {
            // 分支1:正常推理
            result = engine.generate(request.clone(), self.cancel_token.child_token()) => {
                result
            }
            // 分支2:超时
            _ = tokio::time::sleep_until(deadline.into()) => {
                self.handle_timeout(request).await
            }
            // 分支3:上游主动取消
            _ = self.cancel_token.cancelled() => {
                self.handle_cancel(request).await
            }
        }
    }

    async fn handle_timeout(&self, request: Request) -> Result<InferenceResult> {
        match self.cancel_policy {
            CancelPolicy::Immediate => {
                // 发送 engine kill 命令
                Err(InferenceError::Timeout)
            }
            CancelPolicy::CachePartial => {
                // 通知 engine 停止并返回已有的 last_token
                todo!("cache partial generation")
            }
            CancelPolicy::FinishAndCache => {
                // 不设超时,让推理跑完
                todo!("detach and keep running")
            }
        }
    }
}

五、优先级调度:多级队列与 Serverless 场景

5.1 为什么需要优先级

生产推理网关面对的用户通常有明显的 SLA 分层:

用户等级 典型场景 延迟要求
P0 实时交互 Chatbot / Agent p99 < 2s TTFT
P1 商业分析 批量摘要 / RAG p99 < 10s
P2 离线训练 数据清洗 / 标注 尽力而为

如果 P0 和 P2 请求在同一队列中竞争,必须确保 P0 的延迟不被 P2 的长请求污染。

5.2 实现方案:带优先级的多级反馈队列(MLFQ)

use tokio::sync::Semaphore;

/// 优先级队列
struct PriorityQueues {
    /// 每个优先级一个 Semaphore 控制并发
    concurrency_limit: Vec<Semaphore>,
    /// 每个优先级的 FIFO 队列
    queues: Vec<ArrayDeque<PendingRequest>>,
    /// 权重配额(P0 最多占用 60% GPU budget, P1 30%, P2 10%)
    weights: Vec<f64>,
}

impl PriorityQueues {
    fn new() -> Self {
        Self {
            concurrency_limit: vec![
                Semaphore::new(8),  // P0: 最多 8 个并发推理
                Semaphore::new(4),  // P1: 最多 4 个并发推理
                Semaphore::new(2),  // P2: 最多 2 个并发推理
            ],
            queues: vec![ArrayDeque::new(), ArrayDeque::new(), ArrayDeque::new()],
            weights: vec![0.6, 0.3, 0.1],
        }
    }

    /// 从最高优先级开始调度
    async fn dequeue(&mut self) -> Option<PendingRequest> {
        // 严格优先级:P0 非空时总是选 P0
        if !self.queues[0].is_empty() {
            if self.concurrency_limit[0].try_acquire().is_ok() {
                return self.queues[0].pop_front();
            }
        }

        // 加权 P1/P2(公平轮转)
        for priority in 1..=2 {
            if !self.queues[priority].is_empty() {
                if self.concurrency_limit[priority].try_acquire().is_ok() {
                    return self.queues[priority].pop_front();
                }
            }
        }

        None
    }
}

5.3 交互:优先级抢占

最高优先级下,新 P0 请求到达时可以抢占 P2 的推理 slot:

async fn preempt_if_needed(
    task: &PreemptableTask,
    priority: Priority,
    state: &mut GatewayState,
) -> Result<PreemptToken, GatewayError> {
    if priority != Priority::P0 {
        return Err(GatewayError::NotPreemptible);
    }

    // 寻找可抢占的 P2 低优先级任务
    let victim = state.active_tasks.iter_mut()
        .find(|(_, t)| t.priority == Priority::P2 && t.can_preempt);

    if let Some((id, victim)) = victim {
        // 通知被抢占的任务停止生成
        victim.cancel_token.cancel();
        // 将被抢占任务保存的 KV Cache 写入临时存储(用于快速恢复)
        state.preempted_cache.insert(id, victim.capture_cache().await?);

        Ok(PreemptToken { preempted_id: id })
    } else {
        Err(GatewayError::NoPreemptibleTask)
    }
}

六、流式响应与背压控制

6.1 SSE/gRPC Streaming 的 Cancel Safety

推理请求大部分是流式响应(SSE 或 gRPC streaming)。流式响应的一个难题是:生成端的速度可能快于消费端(比如 Agent 在消费 SSE 时需要调用工具、访问数据库),这会导致背压问题。

Rust async 的"零成本"抽象在流控上需要额外注意:

use tokio::sync::mpsc;

fn create_inference_stream(req: Request) -> impl Stream<Item = Result<Token>> {
    let (tx, rx) = mpsc::channel(256); // 缓冲 256 个 token

    tokio::spawn(async move {
        // 启动推理引擎,流式消费 output token stream
        let mut output_stream = engine.generate_stream(req).await;

        while let Some(token) = output_stream.next().await {
            // send 在缓冲区满时挂起(背压传递)
            if tx.send(Ok(token)).await.is_err() {
                // 接收端断开 → 停止推理
                engine.cancel(req.id).await;
                break;
            }
        }
    });

    // 将 mpsc::Receiver 包装为 Stream
    ReceiverStream::new(rx)
}

七、生产实战:Gateway 部署模式

7.1 Sidecar 模式(GPU Node Local)

┌─────────────────────────────────────────┐
│            GPU Node                      │
│  ┌─────────┐  ┌───────────┐  ┌────────┐ │
│  │ Gateway │──│   vLLM    │──│  GPU   │ │
│  │ Sidecar │  │  Engine   │  │ Memory │ │
│  └─────────┘  └───────────┘  └────────┘ │
│        通过 UDS/共享内存通信              │
└─────────────────────────────────────────┘

Gateway 与 Engine 在同一节点上,通过 Unix Domain Socket 通信。适合:KV Cache 亲和性强、需要低延迟 access Engine 状态(GPU Util/KV Cache Stats)。

7.2 Centralized Gateway 模式

                     ┌──────────────┐
                     │  Load Balancer│
                     └──────┬───────┘
                            │
              ┌─────────────┼─────────────┐
              ▼             ▼              ▼
       ┌─────────┐   ┌─────────┐   ┌─────────┐
       │Gateway 1│   │Gateway 2│   │Gateway 3│
       └────┬────┘   └────┬────┘   └────┬────┘
            │              │              │
            ▼              ▼              ▼
       ┌─────────┐   ┌─────────┐   ┌─────────┐
       │ vLLM 1  │   │ vLLM 2  │   │ vLLM 3  │
       └─────────┘   └─────────┘   └─────────┘

Gateway 无状态,根据 etcd/consul 读取 Engine 状态。适合:大规模集群、需要统一接入(认证/限流/审计)的场景。

7.3 推荐拓扑(中大规模)

在中大规模部署(50+ GPU 节点)中,推荐混合模式:

  • 每个 GPU 节点上部署 Gateway Sidecar,负责本地 Engine 状态上报 + 请求去重 + 超时取消
  • 上层 Central Gateway 负责全局路由决策 + 认证 + 计费
  • etcd/consul 作为 Engine 状态的共享存储
/// 简化的混合架构信息流
#[tokio::main]
async fn main() {
    // Sidecar 上报 Engine 状态到 etcd
    let sidecar_handle = tokio::spawn(async {
        let mut engine = VllmEngine::new("0.0.0.0:5555").await;
        let mut ticker = tokio::time::interval(Duration::from_millis(100));

        loop {
            ticker.tick().await;
            let metrics = engine.get_metrics().await;
            etcd_client.put(
                format!("/engines/{}/metrics", engine.id),
                serde_json::to_vec(&metrics).unwrap(),
                Some(PutOpts::new().with_lease(lease_id)),
            ).await.unwrap();
        }
    });

    // Central Gateway 读取 etcd 状态进行路由
    let central_handle = tokio::spawn(async {
        start_central_gateway("0.0.0.0:8080", etcd_client).await;
    });
}

八、Gateway Metrics 体系

8.1 核心 SLI/SLO

/// 推理网关监控指标
struct GatewayMetrics {
    /// 请求延迟分布
    request_latency: HistogramVec,  // [priority, model, status]

    /// KV Cache 命中率
    kv_cache_hit_rate: GaugeVec,   // [model]

    /// 去重效率
    dedup_hit_rate: Gauge,         // 全局

    /// 取消率(按原因分类)
    cancel_rate: CounterVec,       // [reason: timeout/client_disconnect/preempted]

    /// 队列深度
    queue_depth: GaugeVec,         // [priority]

    /// 背压时间(stream send 等待)
    backpressure_duration: Histogram,
}

8.2 SLO 定义示例

SLI SLO 测量方法
P0 请求成功率 ≥ 99.9% 2xx 响应 / 总请求数
P0 TTFT p99 < 2s Prometheus histogram
GPU 闲置率 < 5% 非 idle 时间 / 总时间
取消率 < 1% cancelled / total
去重节省 ≥ 15% QPS in_flight_hit / total

九、总结

生产级 AI 推理网关的核心挑战可以归结为三张表:

路由——让请求去对的地方:感知负载的加权路由 + Prefix-Aware KV Cache 亲和。

去重——让重复请求不浪费资源:进程内 LRU + in-flight Request Coalescing。

取消——让过期请求懂得体面退出:RAII Guard 保证资源回收 + 分层 Cancel Policy。

而优先级调度、流式背压和合理的部署拓扑则是将这些能力串联成稳定生产系统的骨架。

随着多 Agent 协同(Multi-Agent)和长上下文应用的爆发,推理网关层正在成为 AI Infra 中最需要深度工程化的环节之一。

点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论
立即
投稿

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部