AI 推理服务弹性调度与 GPU 资源管理实践

一、场景痛点：GPU 资源稀缺与弹性需求

随着大语言模型在各行业的广泛应用，GPU 资源的管理和调度成为影响 AI 应用性能和成本的核心问题。与传统 CPU 计算不同，GPU 有以下独特挑战：

1. 资源稀缺：高端 GPU（如 H100、A100）供应紧张，价格昂贵
2. 异构资源：不同型号 GPU 性能差异大，需要智能调度
3. 弹性需求波动：AI 推理请求量波动剧烈，需要快速扩缩容
4. 多任务共享：同一 GPU 上可能运行多个推理任务，需要合理分配

传统的资源管理方式（固定分配、手动扩缩容）已经无法满足 AI 时代的挑战，需要更智能的弹性调度方案。

二、底层机制与原理深度剖析

2.1 GPU 资源调度架构

flowchart TD
    subgraph 调度层
        A[API Gateway] --> B[调度器]
        B --> C{调度决策}
    end
    
    subgraph GPU 资源层
        D[GPU Pool Manager]
        E[节点1: A100 x4]
        F[节点2: A100 x4]
        G[节点3: H100 x2]
    end
    
    subgraph AI 推理服务
        H[vLLM Instance 1]
        I[vLLM Instance 2]
        J[TensorRT-LLM Instance]
    end
    
    C -->|资源分配| D
    D --> E
    D --> F
    D --> G
    
    E --> H
    E --> I
    G --> J
    
    K[Prometheus] --> B
    K --> D
    K -->监控指标
    
    style B fill:#b8d4ff
    style D fill:#FFE4B5

调度器的核心职责：

1. 资源感知：了解 GPU 的类型、数量、显存、温度等状态
2. 请求路由：将推理请求路由到合适的 GPU 实例
3. 弹性扩缩：根据负载自动调整实例数量
4. 公平分配：在多个租户之间公平分配 GPU 资源

2.2 GPU 调度算法分类

flowchart LR
    A[调度算法] --> B[基于规则]
    A --> C[基于队列]
    A --> D[基于预测]
    A --> E[基于强化学习]
    
    B --> B1[轮询]
    B --> B2[最少连接]
    B --> B3[亲和性]
    
    C --> C1[优先级队列]
    C --> C2[公平调度]
    C --> C3[资源预留]
    
    D --> D1[流量预测]
    D --> D2[容量规划]
    
    E --> E1[DeepRM]
    E --> E2[Decima]

三、生产级代码实现与最佳实践

3.1 GPU 资源管理器

# ==================== GPU 资源管理器 ====================
"""
生产级 GPU 资源管理系统
支持多节点、多 GPU 的资源调度
"""
import asyncio
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set
from enum import Enum
import threading
import psutil
import subprocess
from collections import defaultdict
import logging

logger = logging.getLogger(__name__)

class GPUState(Enum):
    IDLE = "idle"           # 空闲
    ALLOCATED = "allocated" # 已分配
    RESERVED = "reserved"   # 预留
    FAULTY = "faulty"       # 故障

@dataclass
class GPUInfo:
    """GPU 信息"""
    index: int
    name: str
    memory_total: int       # bytes
    memory_free: int        # bytes
    utilization: float       # 0-1
    temperature: float       # celsius
    state: GPUState = GPUState.IDLE
    allocated_instances: Set[str] = field(default_factory=set)
    
    @property
    def memory_used(self) -> int:
        return self.memory_total - self.memory_free
    
    @property
    def memory_utilization(self) -> float:
        return self.memory_used / self.memory_total if self.memory_total > 0 else 0
    
    @property
    def can_allocate(self) -> bool:
        return self.state == GPUState.IDLE and len(self.allocated_instances) == 0

class GPUResourceManager:
    """
    GPU 资源管理器
    核心功能：
    1. GPU 状态监控
    2. 资源分配与回收
    3. 负载均衡
    4. 故障检测
    """
    
    def __init__(self, config: 'GPUManagerConfig'):
        self.config = config
        self.nodes: Dict[str, List[GPUInfo]] = {}
        self.instance_to_gpu: Dict[str, tuple] = {}  # instance_id -> (node_id, gpu_index)
        self.gpu_to_instance: Dict[tuple, str] = {}  # (node_id, gpu_index) -> instance_id
        
        self._lock = threading.RLock()
        self._monitoring = False
        self._monitor_task: Optional[asyncio.Task] = None
        
    async def start(self):
        """启动资源管理器"""
        await self._discover_gpus()
        self._monitoring = True
        self._monitor_task = asyncio.create_task(self._monitor_loop())
        logger.info("GPU Resource Manager started")
    
    async def stop(self):
        """停止资源管理器"""
        self._monitoring = False
        if self._monitor_task:
            self._monitor_task.cancel()
        logger.info("GPU Resource Manager stopped")
    
    async def _discover_gpus(self):
        """发现集群中的 GPU 资源"""
        # 通过 nvidia-smi 或 NVML 发现 GPU
        try:
            result = subprocess.run(
                ['nvidia-smi', '--query-gpu=index,name,memory.total,memory.free,utilization.gpu,temperature.gpu', '--format=csv,noheader,nounits'],
                capture_output=True,
                text=True,
                check=True
            )
            
            node_id = self._get_node_id()
            
            for line in result.stdout.strip().split('\n'):
                if not line:
                    continue
                    
                parts = [p.strip() for p in line.split(',')]
                if len(parts) != 6:
                    continue
                
                gpu_info = GPUInfo(
                    index=int(parts[0]),
                    name=parts[1],
                    memory_total=int(parts[2]) * 1024 * 1024,  # MB to bytes
                    memory_free=int(parts[3]) * 1024 * 1024,
                    utilization=float(parts[4]) / 100,
                    temperature=float(parts[5])
                )
                
                if node_id not in self.nodes:
                    self.nodes[node_id] = []
                self.nodes[node_id].append(gpu_info)
                
            logger.info(f"Discovered {len(self.nodes.get(node_id, []))} GPUs on node {node_id}")
            
        except Exception as e:
            logger.warning(f"Failed to discover GPUs: {e}")
    
    def _get_node_id(self) -> str:
        """获取节点 ID"""
        return subprocess.run(['hostname'], capture_output=True, text=True).stdout.strip()
    
    async def _monitor_loop(self):
        """监控循环"""
        while self._monitoring:
            try:
                await self._update_gpu_status()
                await self._check_gpu_health()
                await asyncio.sleep(self.config.monitoring_interval)
            except Exception as e:
                logger.error(f"Monitoring error: {e}")
    
    async def _update_gpu_status(self):
        """更新 GPU 状态"""
        try:
            result = subprocess.run(
                ['nvidia-smi', '--query-gpu=index,memory.free,utilization.gpu,temperature.gpu', '--format=csv,noheader,nounits'],
                capture_output=True,
                text=True,
                check=True
            )
            
            node_id = self._get_node_id()
            gpus = self.nodes.get(node_id, [])
            
            for line, gpu in zip(result.stdout.strip().split('\n'), gpus):
                if not line:
                    continue
                    
                parts = [p.strip() for p in line.split(',')]
                if len(parts) != 4:
                    continue
                
                with self._lock:
                    gpu.memory_free = int(parts[1]) * 1024 * 1024
                    gpu.utilization = float(parts[2]) / 100
                    gpu.temperature = float(parts[3])
                    
        except Exception as e:
            logger.warning(f"Failed to update GPU status: {e}")
    
    async def _check_gpu_health(self):
        """检查 GPU 健康状态"""
        node_id = self._get_node_id()
        gpus = self.nodes.get(node_id, [])
        
        for gpu in gpus:
            with self._lock:
                # 温度过高的 GPU 标记为故障
                if gpu.temperature > self.config.max_temperature:
                    if gpu.state != GPUState.FAULTY:
                        logger.warning(f"GPU {gpu.index} temperature too high: {gpu.temperature}°C")
                        gpu.state = GPUState.FAULTY
                        
                # 显存不足的 GPU 标记为 Reserve
                if gpu.memory_free < self.config.min_free_memory:
                    if gpu.state != GPUState.RESERVED:
                        logger.warning(f"GPU {gpu.index} low memory: {gpu.memory_free / 1024**2:.0f}MB free")
                        gpu.state = GPUState.RESERVED
    
    def allocate_gpu(
        self,
        instance_id: str,
        memory_required: int,
        preference: Optional[Dict] = None
    ) -> Optional[tuple]:
        """
        分配 GPU 资源
        返回 (node_id, gpu_index) 或 None
        """
        with self._lock:
            for node_id, gpus in self.nodes.items():
                for gpu in sorted(gpus, key=lambda g: g.memory_free, reverse=True):
                    # 检查是否可用
                    if gpu.state == GPUState.FAULTY:
                        continue
                    if gpu.memory_free < memory_required:
                        continue
                    if gpu.state == GPUState.RESERVED and (
                        preference is None or not preference.get('allow_reserved', False)
                    ):
                        continue
                    
                    # 分配 GPU
                    gpu.state = GPUState.ALLOCATED
                    gpu.allocated_instances.add(instance_id)
                    
                    self.instance_to_gpu[instance_id] = (node_id, gpu.index)
                    self.gpu_to_instance[(node_id, gpu.index)] = instance_id
                    
                    logger.info(f"Allocated GPU {node_id}:{gpu.index} to instance {instance_id}")
                    return (node_id, gpu.index)
            
            return None
    
    def release_gpu(self, instance_id: str) -> bool:
        """释放 GPU 资源"""
        with self._lock:
            if instance_id not in self.instance_to_gpu:
                logger.warning(f"Instance {instance_id} not allocated")
                return False
            
            node_id, gpu_index = self.instance_to_gpu[instance_id]
            
            if node_id in self.nodes:
                for gpu in self.nodes[node_id]:
                    if gpu.index == gpu_index:
                        gpu.allocated_instances.discard(instance_id)
                        if len(gpu.allocated_instances) == 0:
                            gpu.state = GPUState.IDLE
                        logger.info(f"Released GPU {node_id}:{gpu.index} from instance {instance_id}")
                        break
            
            del self.instance_to_gpu[instance_id]
            del self.gpu_to_instance[(node_id, gpu_index)]
            
            return True
    
    def get_allocation_summary(self) -> Dict:
        """获取资源分配摘要"""
        with self._lock:
            total_gpus = sum(len(gpus) for gpus in self.nodes.values())
            allocated_gpus = sum(
                1 for gpus in self.nodes.values() 
                for gpu in gpus if gpu.state == GPUState.ALLOCATED
            )
            
            return {
                'total_gpus': total_gpus,
                'allocated_gpus': allocated_gpus,
                'idle_gpus': total_gpus - allocated_gpus,
                'utilization': allocated_gpus / total_gpus if total_gpus > 0 else 0,
                'by_node': {
                    node_id: {
                        'total': len(gpus),
                        'allocated': sum(1 for g in gpus if g.state == GPUState.ALLOCATED),
                        'idle': sum(1 for g in gpus if g.state == GPUState.IDLE),
                    }
                    for node_id, gpus in self.nodes.items()
                }
            }

@dataclass
class GPUManagerConfig:
    """资源配置"""
    monitoring_interval: int = 5  # 秒
    max_temperature: float = 85.0  # celsius
    min_free_memory: int = 2 * 1024 * 1024 * 1024  # 2GB

3.2 弹性调度器

# ==================== AI 推理弹性调度器 ====================
"""
基于预测的弹性调度器
支持：
1. 主动扩缩容
2. 流量预测
3. 蓝绿部署
4. 金丝雀发布
"""
import asyncio
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set
from collections import deque
import logging
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

@dataclass
class ScalingConfig:
    """扩缩容配置"""
    min_instances: int = 1
    max_instances: int = 10
    scale_up_threshold: float = 0.8  # 80% 利用率触发扩容
    scale_down_threshold: float = 0.3  # 30% 利用率触发缩容
    scale_up_cooldown: int = 60  # 扩容冷却时间（秒）
    scale_down_cooldown: int = 300  # 缩容冷却时间（秒）
    target_utilization: float = 0.7  # 目标利用率

@dataclass
class InstanceInfo:
    """推理实例信息"""
    instance_id: str
    gpu_node: str
    gpu_index: int
    model_name: str
    memory_required: int
    current_load: float = 0.0
    request_count: int = 0
    created_at: datetime = field(default_factory=datetime.now)
    status: str = "starting"  # starting, ready, draining, stopped

class ElasticScheduler:
    """
    弹性调度器
    核心功能：
    1. 基于利用率的自动扩缩容
    2. 基于时间序列的预测性扩容
    3. 流量分配与负载均衡
    """
    
    def __init__(
        self,
        resource_manager: 'GPUResourceManager',
        scaling_config: ScalingConfig
    ):
        self.resource_manager = resource_manager
        self.scaling_config = scaling_config
        
        self.instances: Dict[str, InstanceInfo] = {}
        self.instances_by_model: Dict[str, Set[str]] = {}
        
        # 流量历史（用于预测）
        self.request_history: deque = deque(maxlen=1000)
        
        # 指标收集
        self.metrics_history: deque = deque(maxlen=100)
        
        # 扩缩容状态
        self.last_scale_up_time: datetime = datetime.min
        self.last_scale_down_time: datetime = datetime.min
        
        self._scheduler_task: Optional[asyncio.Task] = None
        self._running = False
    
    async def start(self):
        """启动调度器"""
        self._running = True
        self._scheduler_task = asyncio.create_task(self._scheduler_loop())
        logger.info("Elastic Scheduler started")
    
    async def stop(self):
        """停止调度器"""
        self._running = False
        if self._scheduler_task:
            self._scheduler_task.cancel()
        logger.info("Elastic Scheduler stopped")
    
    async def _scheduler_loop(self):
        """调度循环"""
        while self._running:
            try:
                # 收集指标
                await self._collect_metrics()
                
                # 预测流量
                predicted_load = self._predict_load()
                
                # 执行扩缩容决策
                await self._execute_scaling_decision(predicted_load)
                
                # 负载均衡
                await self._rebalance_load()
                
                await asyncio.sleep(10)  # 每 10 秒检查一次
                
            except Exception as e:
                logger.error(f"Scheduler error: {e}")
    
    async def _collect_metrics(self):
        """收集实例指标"""
        total_load = 0.0
        total_instances = 0
        
        for instance in self.instances.values():
            if instance.status == "ready":
                # 模拟指标收集（实际应从实例获取）
                instance.current_load = min(1.0, instance.request_count / 100)
                total_load += instance.current_load
                total_instances += 1
        
        if total_instances > 0:
            avg_utilization = total_load / total_instances
            self.metrics_history.append({
                'timestamp': datetime.now(),
                'utilization': avg_utilization,
                'instances': total_instances
            })
    
    def _predict_load(self) -> float:
        """基于时间序列预测未来负载"""
        if len(self.request_history) < 10:
            return 0.5  # 默认 50%
        
        # 简单移动平均预测
        recent = list(self.request_history)[-30:]
        if not recent:
            return 0.5
        
        avg_load = sum(r['load'] for r in recent) / len(recent)
        
        # 检测趋势
        if len(recent) >= 60:
            older = sum(r['load'] for r in recent[-60:-30]) / 30
            newer = sum(r['load'] for r in recent[-30:]) / 30
            trend = (newer - older) / (older + 1e-6)
        else:
            trend = 0
        
        # 预测
        predicted = avg_load * (1 + trend * 0.5)
        return max(0, min(1, predicted))
    
    async def _execute_scaling_decision(self, predicted_load: float):
        """执行扩缩容决策"""
        now = datetime.now()
        current_instances = sum(1 for i in self.instances.values() if i.status == "ready")
        
        # 计算目标实例数
        target_instances = current_instances
        
        # 获取当前平均利用率
        current_utilization = 0
        if self.metrics_history:
            current_utilization = self.metrics_history[-1]['utilization']
        
        # 扩容决策
        if current_utilization > self.scaling_config.scale_up_threshold:
            if (now - self.last_scale_up_time).total_seconds() > self.scaling_config.scale_up_cooldown:
                if current_instances < self.scaling_config.max_instances:
                    target_instances = min(
                        self.scaling_config.max_instances,
                        current_instances + 1
                    )
                    self.last_scale_up_time = now
                    logger.info(f"Scaling up: {current_instances} -> {target_instances}")
        
        # 缩容决策
        elif current_utilization < self.scaling_config.scale_down_threshold:
            if (now - self.last_scale_down_time).total_seconds() > self.scaling_config.scale_down_cooldown:
                if current_instances > self.scaling_config.min_instances:
                    target_instances = max(
                        self.scaling_config.min_instances,
                        current_instances - 1
                    )
                    self.last_scale_down_time = now
                    logger.info(f"Scaling down: {current_instances} -> {target_instances}")
        
        # 执行扩缩容
        if target_instances != current_instances:
            if target_instances > current_instances:
                await self._scale_up(target_instances - current_instances)
            else:
                await self._scale_down(current_instances - target_instances)
    
    async def _scale_up(self, count: int):
        """扩容"""
        for _ in range(count):
            instance_id = f"inst_{int(time.time() * 1000)}"
            
            # 分配 GPU
            gpu_allocation = self.resource_manager.allocate_gpu(
                instance_id=instance_id,
                memory_required=8 * 1024**3  # 8GB
            )
            
            if gpu_allocation is None:
                logger.warning(f"Failed to allocate GPU for new instance")
                break
            
            node_id, gpu_index = gpu_allocation
            
            instance = InstanceInfo(
                instance_id=instance_id,
                gpu_node=node_id,
                gpu_index=gpu_index,
                model_name="default",
                memory_required=8 * 1024**3
            )
            
            self.instances[instance_id] = instance
            
            if instance.model_name not in self.instances_by_model:
                self.instances_by_model[instance.model_name] = set()
            self.instances_by_model[instance.model_name].add(instance_id)
            
            # 异步启动实例
            asyncio.create_task(self._start_instance(instance))
    
    async def _scale_down(self, count: int):
        """缩容"""
        # 选择最空闲的实例
        sorted_instances = sorted(
            [i for i in self.instances.values() if i.status == "ready"],
            key=lambda x: x.current_load
        )
        
        for instance in sorted_instances[:count]:
            await self._stop_instance(instance)
    
    async def _start_instance(self, instance: InstanceInfo):
        """启动推理实例"""
        instance.status = "starting"
        
        # 模拟启动过程（实际应启动 vLLM 等）
        await asyncio.sleep(5)
        
        instance.status = "ready"
        logger.info(f"Instance {instance.instance_id} started on {instance.gpu_node}:{instance.gpu_index}")
    
    async def _stop_instance(self, instance: InstanceInfo):
        """停止推理实例"""
        instance.status = "draining"
        
        # 等待现有请求处理完成
        await asyncio.sleep(10)
        
        # 释放 GPU
        self.resource_manager.release_gpu(instance.instance_id)
        
        # 移除实例
        self.instances_by_model[instance.model_name].discard(instance.instance_id)
        del self.instances[instance.instance_id]
        
        logger.info(f"Instance {instance.instance_id} stopped")
    
    async def _rebalance_load(self):
        """负载均衡"""
        if len(self.request_history) == 0:
            return
        
        # 获取最新请求的模型
        latest_request = self.request_history[-1]
        model_name = latest_request.get('model', 'default')
        
        # 选择负载最低的实例
        ready_instances = [
            i for i in self.instances.values() 
            if i.status == "ready" and i.model_name == model_name
        ]
        
        if not ready_instances:
            return
        
        # 选择最空闲的实例
        selected = min(ready_instances, key=lambda x: x.current_load)
        
        logger.debug(f"Selected instance {selected.instance_id} with load {selected.current_load}")
    
    async def route_request(
        self,
        model_name: str,
        request_data: dict
    ) -> Optional[str]:
        """路由请求到合适实例"""
        self.request_history.append({
            'timestamp': datetime.now(),
            'model': model_name,
            'load': 0.5  # 简化
        })
        
        ready_instances = [
            i for i in self.instances.values() 
            if i.status == "ready" and i.model_name == model_name
        ]
        
        if not ready_instances:
            return None
        
        # 简单轮询
        return ready_instances[0].instance_id
    
    def get_status(self) -> dict:
        """获取调度器状态"""
        return {
            'total_instances': len(self.instances),
            'ready_instances': sum(1 for i in self.instances.values() if i.status == "ready"),
            'metrics': {
                'avg_utilization': (
                    self.metrics_history[-1]['utilization'] 
                    if self.metrics_history else 0
                ),
                'request_count': len(self.request_history),
            },
            'scaling': {
                'last_scale_up': self.last_scale_up_time.isoformat(),
                'last_scale_down': self.last_scale_down_time.isoformat(),
            }
        }

四、边界分析与架构权衡

4.1 GPU 调度策略对比

4.2 弹性调度注意事项

五、总结

AI 推理服务的弹性调度是 AI 基础设施的核心能力。通过智能的资源管理和调度算法，可以实现：

1. 资源高效利用：最大化 GPU 利用率，降低单位推理成本
2. 弹性伸缩：快速响应流量变化，保证服务质量
3. 成本优化：预测性扩容，避免资源浪费
4. 高可用：故障自动检测和恢复，保证服务稳定性

关键成功因素：

1. 完善的监控：实时了解 GPU 状态和负载情况
2. 智能调度算法：结合预测和实时状态做决策
3. 资源预留：为突发流量预留缓冲资源
4. 渐进式实施：从简单策略开始，逐步引入 AI

GPU 资源管理的智能化是 AI 时代运维的核心挑战，需要持续投入和优化。