模型推理部署实战:从GPU资源调度到推理服务化的全链路方案

cover

一、推理部署的现实困境:模型跑得动,服务扛不住

大模型推理部署是AI应用落地的最后一公里,也是最容易被低估的环节。训练一个模型可能需要几周,但把它部署成稳定可用的推理服务,同样需要大量工程工作。

最常见的困境:模型在开发环境跑得好好的,一上线就出问题。GPU显存不够导致OOM,并发请求一多延迟就飙升,模型加载时间太长影响服务启动,多模型共存时资源争抢严重。这些问题不是模型本身的问题,而是推理服务化的工程问题。

更深层的问题是成本。GPU资源昂贵,一个A100实例月租上万元。如果推理服务的GPU利用率只有20%,意味着80%的钱在浪费。但提高利用率又面临风险:并发太高可能导致显存溢出,批处理太大会增加延迟。如何在成本和性能之间找到最优解,是推理部署的核心挑战。

二、推理服务化的核心架构

2.1 推理服务整体架构

flowchart TB
    subgraph "接入层"
        A[API Gateway] --> B[负载均衡]
        B --> C[请求队列]
    end

    subgraph "调度层"
        C --> D[推理调度器]
        D --> E[批处理管理器]
        D --> F[模型路由器]
    end

    subgraph "推理层"
        F --> G1[模型实例A<br/>GPU:0]
        F --> G2[模型实例B<br/>GPU:1]
        F --> G3[模型实例C<br/>CPU]
    end

    subgraph "运维层"
        H[GPU监控]
        I[自动扩缩容]
        J[模型热加载]
        K[指标采集]
    end

    G1 --> H
    G2 --> H
    H --> I
    J --> F
    G1 --> K
    G2 --> K
    G3 --> K

2.2 推理请求生命周期

一个推理请求从到达API到返回结果,经历以下阶段:请求解析与校验 → Token计数与限流 → 队列等待 → 批处理聚合 → GPU推理 → 后处理 → 响应返回。其中,批处理聚合和GPU推理是性能关键路径。

三、生产级推理服务实现

3.1 推理调度器

package inference

import (
    "container/heap"
    "context"
    "sync"
    "time"
)

// InferenceRequest 推理请求
type InferenceRequest struct {
    ID        string
    ModelName string
    Input     string
    MaxTokens int
    Priority  int       // 优先级,数值越小越优先
    SubmitAt  time.Time // 提交时间
    ResultCh  chan *InferenceResult
}

// InferenceResult 推理结果
type InferenceResult struct {
    Output    string
    Tokens    int
    Latency   time.Duration
    Error     error
}

// PriorityQueue 优先级队列
type PriorityQueue []*InferenceRequest

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
    // 优先级相同时,先提交的优先
    if pq[i].Priority == pq[j].Priority {
        return pq[i].SubmitAt.Before(pq[j].SubmitAt)
    }
    return pq[i].Priority < pq[j].Priority
}

func (pq PriorityQueue) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
}

func (pq *PriorityQueue) Push(x interface{}) {
    item := x.(*InferenceRequest)
    *pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    *pq = old[:n-1]
    return item
}

// InferenceScheduler 推理调度器
type InferenceScheduler struct {
    mu       sync.Mutex
    queue    PriorityQueue
    models   map[string]*ModelInstance
    maxBatch int           // 最大批处理大小
    maxWait  time.Duration // 最大等待时间
}

// ModelInstance 模型实例
type ModelInstance struct {
    Name       string
    Device     string // gpu:0, cpu
    MaxConcurrent int
    currentLoad   int
    isReady       bool
}

// NewInferenceScheduler 创建调度器
func NewInferenceScheduler(maxBatch int, maxWait time.Duration) *InferenceScheduler {
    return &InferenceScheduler{
        queue:    make(PriorityQueue, 0),
        models:   make(map[string]*ModelInstance),
        maxBatch: maxBatch,
        maxWait:  maxWait,
    }
}

// RegisterModel 注册模型实例
func (s *InferenceScheduler) RegisterModel(name, device string, maxConcurrent int) {
    s.mu.Lock()
    defer s.mu.Unlock()
    s.models[name] = &ModelInstance{
        Name:          name,
        Device:        device,
        MaxConcurrent: maxConcurrent,
        isReady:       true,
    }
}

// Submit 提交推理请求
func (s *InferenceScheduler) Submit(ctx context.Context, req *InferenceRequest) (*InferenceResult, error) {
    req.SubmitAt = time.Now()
    req.ResultCh = make(chan *InferenceResult, 1)

    s.mu.Lock()
    heap.Push(&s.queue, req)
    s.mu.Unlock()

    select {
    case result := <-req.ResultCh:
        return result, nil
    case <-ctx.Done():
        return nil, ctx.Err()
    }
}

// ScheduleLoop 调度循环
func (s *InferenceScheduler) ScheduleLoop(ctx context.Context) {
    ticker := time.NewTicker(10 * time.Millisecond)
    defer ticker.Stop()

    for {
        select {
        case <-ctx.Done():
            return
        case <-ticker.C:
            s.scheduleBatch()
        }
    }
}

// scheduleBatch 调度一批请求
func (s *InferenceScheduler) scheduleBatch() {
    s.mu.Lock()
    defer s.mu.Unlock()

    if s.queue.Len() == 0 {
        return
    }

    // 按模型分组收集请求
    batches := make(map[string][]*InferenceRequest)

    for s.queue.Len() > 0 {
        req := heap.Pop(&s.queue).(*InferenceRequest)

        // 检查请求是否超时
        if time.Since(req.SubmitAt) > s.maxWait*2 {
            req.ResultCh <- &InferenceResult{
                Error: fmt.Errorf("request timeout in queue"),
            }
            continue
        }

        model, ok := s.models[req.ModelName]
        if !ok || !model.isReady {
            req.ResultCh <- &InferenceResult{
                Error: fmt.Errorf("model %s not available", req.ModelName),
            }
            continue
        }

        batches[req.ModelName] = append(batches[req.ModelName], req)

        // 达到最大批处理大小,停止收集
        if len(batches[req.ModelName]) >= s.maxBatch {
            break
        }
    }

    // 分发批处理任务
    for modelName, batch := range batches {
        model := s.models[modelName]
        if model.currentLoad >= model.MaxConcurrent {
            // 模型负载已满,重新入队
            for _, req := range batch {
                heap.Push(&s.queue, req)
            }
            continue
        }

        go s.executeBatch(modelName, batch)
    }
}

// executeBatch 执行批处理推理
func (s *InferenceScheduler) executeBatch(modelName string, batch []*InferenceRequest) {
    model := s.models[modelName]
    model.currentLoad++
    defer func() { model.currentLoad-- }()

    // 构建批处理输入
    inputs := make([]string, len(batch))
    for i, req := range batch {
        inputs[i] = req.Input
    }

    start := time.Now()
    outputs, err := s.callModel(modelName, inputs)
    latency := time.Since(start)

    for i, req := range batch {
        if err != nil {
            req.ResultCh <- &InferenceResult{Error: err}
        } else if i < len(outputs) {
            req.ResultCh <- &InferenceResult{
                Output:  outputs[i],
                Latency: latency,
            }
        } else {
            req.ResultCh <- &InferenceResult{
                Error: fmt.Errorf("batch result index out of range"),
            }
        }
    }
}

3.2 GPU资源管理器

package inference

import (
    "context"
    "fmt"
    "sync"
    "time"
)

// GPUInfo GPU信息
type GPUInfo struct {
    Index       int
    Name        string
    MemoryTotal uint64 // 总显存MB
    MemoryUsed  uint64 // 已用显存MB
    Utilization int    // GPU利用率%
    Temperature int    // 温度℃
}

// GPUManager GPU资源管理器
type GPUManager struct {
    mu       sync.RWMutex
    gpus     map[int]*GPUInfo
    allocate map[int][]string // GPU索引 -> 已分配的模型名列表
}

func NewGPUManager() *GPUManager {
    return &GPUManager{
        gpus:     make(map[int]*GPUInfo),
        allocate: make(map[int][]string),
    }
}

// AllocateGPU 为模型分配GPU
func (gm *GPUManager) AllocateGPU(modelName string, requiredMemoryMB uint64) (int, error) {
    gm.mu.Lock()
    defer gm.mu.Unlock()

    // 查找有足够显存的GPU
    for idx, gpu := range gm.gpus {
        available := gpu.MemoryTotal - gpu.MemoryUsed
        if available >= requiredMemoryMB {
            gpu.MemoryUsed += requiredMemoryMB
            gm.allocate[idx] = append(gm.allocate[idx], modelName)
            return idx, nil
        }
    }

    return -1, fmt.Errorf("no GPU with enough memory: required %dMB", requiredMemoryMB)
}

// ReleaseGPU 释放GPU资源
func (gm *GPUManager) ReleaseGPU(gpuIndex int, modelName string, memoryMB uint64) {
    gm.mu.Lock()
    defer gm.mu.Unlock()

    gpu, ok := gm.gpus[gpuIndex]
    if !ok {
        return
    }

    gpu.MemoryUsed -= memoryMB
    models := gm.allocate[gpuIndex]
    for i, name := range models {
        if name == modelName {
            gm.allocate[gpuIndex] = append(models[:i], models[i+1:]...)
            break
        }
    }
}

// GetAvailableMemory 获取可用显存
func (gm *GPUManager) GetAvailableMemory(gpuIndex int) uint64 {
    gm.mu.RLock()
    defer gm.mu.RUnlock()

    if gpu, ok := gm.gpus[gpuIndex]; ok {
        return gpu.MemoryTotal - gpu.MemoryUsed
    }
    return 0
}

// MonitorLoop GPU监控循环
func (gm *GPUManager) MonitorLoop(ctx context.Context) {
    ticker := time.NewTicker(5 * time.Second)
    defer ticker.Stop()

    for {
        select {
        case <-ctx.Done():
            return
        case <-ticker.C:
            gm.collectGPUMetrics()
        }
    }
}

// collectGPUMetrics 采集GPU指标
func (gm *GPUManager) collectGPUMetrics() {
    // 实际实现中调用nvidia-smi或NVML库
    // 这里是示意代码
}

3.3 模型热加载

package inference

import (
    "sync"
    "sync/atomic"
)

// ModelRegistry 模型注册表
type ModelRegistry struct {
    mu      sync.RWMutex
    models  map[string]*ModelHandle
    loading map[string]bool // 正在加载的模型
}

// ModelHandle 模型句柄
type ModelHandle struct {
    Name         string
    Version      string
    Device       string
    MemoryUsage  uint64
    Status       int32 // 0=loading, 1=ready, 2=unloading
    LoadTime     time.Time
    RequestCount int64
}

// NewModelRegistry 创建模型注册表
func NewModelRegistry() *ModelRegistry {
    return &ModelRegistry{
        models:  make(map[string]*ModelHandle),
        loading: make(map[string]bool),
    }
}

// LoadModel 加载模型(支持热加载)
func (mr *ModelRegistry) LoadModel(name, version, device string) error {
    mr.mu.Lock()
    // 检查是否正在加载
    if mr.loading[name] {
        mr.mu.Unlock()
        return fmt.Errorf("model %s is already loading", name)
    }
    mr.loading[name] = true
    mr.mu.Unlock()

    defer func() {
        mr.mu.Lock()
        delete(mr.loading, name)
        mr.mu.Unlock()
    }()

    // 创建模型句柄
    handle := &ModelHandle{
        Name:     name,
        Version:  version,
        Device:   device,
        Status:   0, // loading
        LoadTime: time.Now(),
    }

    mr.mu.Lock()
    mr.models[name] = handle
    mr.mu.Unlock()

    // 执行模型加载(耗时操作)
    if err := mr.doLoadModel(name, version, device); err != nil {
        mr.mu.Lock()
        delete(mr.models, name)
        mr.mu.Unlock()
        return err
    }

    // 加载完成,标记为ready
    atomic.StoreInt32(&handle.Status, 1)
    return nil
}

// UnloadModel 卸载模型
func (mr *ModelRegistry) UnloadModel(name string) error {
    mr.mu.RLock()
    handle, ok := mr.models[name]
    if !ok {
        mr.mu.RUnlock()
        return fmt.Errorf("model %s not found", name)
    }
    mr.mu.RUnlock()

    // 标记为卸载中,拒绝新请求
    atomic.StoreInt32(&handle.Status, 2)

    // 等待进行中的请求完成
    for atomic.LoadInt64(&handle.RequestCount) > 0 {
        time.Sleep(100 * time.Millisecond)
    }

    // 执行卸载
    if err := mr.doUnloadModel(name); err != nil {
        return err
    }

    mr.mu.Lock()
    delete(mr.models, name)
    mr.mu.Unlock()

    return nil
}

// GetModel 获取可用模型
func (mr *ModelRegistry) GetModel(name string) (*ModelHandle, error) {
    mr.mu.RLock()
    defer mr.mu.RUnlock()

    handle, ok := mr.models[name]
    if !ok {
        return nil, fmt.Errorf("model %s not found", name)
    }

    if atomic.LoadInt32(&handle.Status) != 1 {
        return nil, fmt.Errorf("model %s is not ready", name)
    }

    atomic.AddInt64(&handle.RequestCount, 1)
    return handle, nil
}

// ReleaseModel 释放模型引用
func (mr *ModelRegistry) ReleaseModel(name string) {
    mr.mu.RLock()
    handle, ok := mr.models[name]
    mr.mu.RUnlock()

    if ok {
        atomic.AddInt64(&handle.RequestCount, -1)
    }
}

func (mr *ModelRegistry) doLoadModel(name, version, device string) error {
    // 实际实现:调用模型加载框架(如vLLM、TGI)
    return nil
}

func (mr *ModelRegistry) doUnloadModel(name string) error {
    // 实际实现:释放模型占用的GPU显存
    return nil
}

四、推理部署的边界与权衡

4.1 批处理大小与延迟

批处理可以提高GPU利用率,但会增加请求延迟。等待凑批的时间是额外开销。对于延迟敏感的场景(如实时对话),批处理大小不应超过4-8;对于吞吐优先的场景(如批量处理),可以设置更大的批处理。

4.2 GPU显存与模型大小

大模型的显存占用远超预期。一个7B参数的模型,FP16精度下需要14GB显存,加上KV Cache和运行时开销,实际需要20-25GB。一个80GB的A100只能同时服务2-3个7B模型。量化(INT8/INT4)可以减少显存占用,但会损失精度。

4.3 冷启动问题

模型首次加载需要从磁盘读取权重到GPU,耗时可能长达数十秒。对于弹性扩缩容场景,冷启动延迟可能导致服务短时间不可用。解决方案:保持一定数量的预热实例,或在低峰期预加载模型。

4.4 禁用场景

GPU推理不适合以下场景:请求量极低(GPU闲置浪费);模型极小(CPU推理即可满足);需要严格确定性输出(GPU浮点运算有微小差异);成本极度敏感(CPU推理成本远低于GPU)。

五、总结

模型推理部署是将AI能力转化为生产服务的关键环节。推理调度器管理请求队列和批处理,GPU资源管理器分配和监控显存,模型热加载支持不停机更新。三者协同,构成推理服务化的基础设施。

工程落地的核心考量:批处理大小需要根据延迟SLA和吞吐目标权衡;GPU显存分配需要预留KV Cache和运行时开销;冷启动问题需要预热实例或预加载策略;量化是降低显存和成本的有效手段,但需要评估精度损失。

推理部署不是一次性工作,而是持续优化的过程。监控GPU利用率、推理延迟、请求排队时间等指标,根据数据不断调整调度策略和资源配置。

Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐