AI for Systems:机器学习驱动的系统优化工程实践

当传统系统优化逼近理论极限,机器学习正在重新定义数据库、操作系统和网络基础设施的性能天花板。从 Google 的 learned index 到阿里云的 AI 驱动查询优化器,"AI for Systems" 正在从学术概念走向大规模生产部署。


一、引言:从 Systems for AI 到 AI for Systems

过去十年,系统社区的核心叙事是"为 AI 设计更好的系统"——GPU 集群、分布式训练框架、推理加速器。但一个反向的趋势正在兴起:用 AI 来优化系统本身。

这个方向被称为 AI for Systems,核心思想是利用机器学习模型替代或增强传统系统中基于规则和启发式的组件:

传统组件 学习替代品 代表工作
B+Tree 索引 Learned Index (Kraska et al., 2018) Google Bigtable, SageDB
LRU 缓存替换 Learned Cache (RL-based) Meta CDN, 阿里云 Tair
查询优化器 Neo / Bao / CardBench Microsoft, 华为 GaussDB
CPU 调度器 AI Scheduler (DeepRM) Google Borg, 阿里 Sigma
拥塞控制 Remy / Aurora / Orca Google BBRv2 探索
哈希函数 Morton Filter / Learned Hash Facebook Zstandard

本文深入剖析这些学习增强型系统的架构设计、训练方法和生产部署挑战。


二、学习索引:从 B+Tree 到递归模型索引

2.1 核心思想

Kraska 等人在 2018 年提出的 Learned Index 洞察了一个简单事实:B+Tree 本质上是一个从 key 到 position 的函数映射。如果用神经网络来学习这个映射,理论上可以达到 O(1) 的查询复杂度。

``` 传统 B+Tree: key → [tree traversal] → page → [binary search] → position Learned Index: key → [neural network inference] → position → [local search] → exact ```

2.2 递归模型索引(RMI)

单层神经网络精度有限,实践中使用 递归模型索引(Recursive Model Index):

```python class RMI: """两阶段递归模型索引""" def __init__(self, stage1_models, stage2_models): self.stage1 = stage1_models # 顶层路由模型 self.stage2 = stage2_models # 底层精确模型 def predict(self, key): # Stage 1: 选择专家模型 expert_id = self.stage1.predict(key) # Stage 2: 精确位置预测 position = self.stage2[expert_id].predict(key) return position def search(self, key, error_bound): """在预测位置 ±error_bound 范围内二分查找""" pos = self.predict(key) low = max(0, pos - error_bound) high = min(len(data), pos + error_bound) return binary_search(data, key, low, high) ```

在实际生产实现中(如 Google 的 SageDB),RMI 的关键设计参数是:

  • 误差上界 (Error Bound): 全局 vs 局部误差,直接影响搜索区间大小
  • 模型复杂度: 线性模型 vs 小型神经网络(通常 1-2 层 MLP,hidden_size ≤ 32)
  • 数据分布变化: 需要在线再训练机制

2.3 ALEX:对抗写入的进化索引

RMI 的致命弱点是对写入不友好。ALEX (Adaptive Learned Index) 解决了这个问题:

```python class ALEXNode: """ALEX 的节点结构 - 结合 B+Tree 的分裂能力和学习模型的预测能力""" def __init__(self, capacity): self.keys = [] # 存储的 key self.model = None # 节点内的线性模型 self.gaps = [] # key 间间隙的统计模型 self.overflow_pages = [] # 溢出页链表 def insert(self, key, value): predicted_pos = self.model.predict(key) # 如果预测位置有空位,直接插入 if self.has_space_at(predicted_pos): self.insert_at(predicted_pos, key, value) else: # 线性 probing 或分裂节点 self.probe_and_insert(key, value) if self.is_overflowing(): self.split() def split(self): """基于模型预测密度的非均匀分裂""" # ALEX 不是均匀分裂,而是根据模型预测的密度分布分裂 # 数据密集区域更小,稀疏区域更大 split_point = self.find_density_inflection_point() return self.do_split(split_point) ```

ALEX 在 YCSB 基准测试中相比 B+Tree 提升 1.5-2.8x 的读写吞吐,同时保持了对数据分布的适应性。

2.4 生产部署的鸿沟

然而,真正的生产部署中存在几个被学术论文忽视的挑战:

  1. 训练成本 vs 收益: 对于 10 亿条记录,训练一个 RMI 需要几分钟到几小时,这仅在数据分布稳定时划算
  2. 并发控制: 学习索引的更新需要模型再训练,与 MVCC 的交互复杂
  3. 内存碎片: 模型的预测误差导致不均匀的存储布局,影响内存友好性
  4. 冷启动: 新索引需要全部数据来训练初始模型
  5. Google 在 Bigtable 中的实际经验是:learned index 在只读或读多写少的场景下效果最佳;高频写入场景下,维护模型的 overhead 可能抵消收益。


    三、学习缓存替换:超越 LRU

    3.1 LRU 的理论局限

    LRU 假设"最近使用的数据最可能再被使用",但在现代工作负载下这个假设经常失效:

    • 顺序扫描污染: 大数据分析的一次性扫描会驱逐热数据
    • 访问频率倾斜: 少量热点数据被频繁访问,大量数据仅访问一次
    • 对象大小不均: LRU 不考虑对象大小,导致大对象占用过多缓存空间

    3.2 Learned Cache:基于 RL 的替换策略

    Meta 在 CDN 和内容分发场景中探索了基于强化学习的缓存替换:

    ```python class LearnedCache: """基于 RL 的缓存替换核心逻辑""" def __init__(self, capacity, state_dim): self.capacity = capacity self.cache = {} # 状态特征:对象大小、频率、时间局部性、内容类型等 self.policy_net = PolicyNetwork(state_dim, hidden=64) self.replay_buffer = ReplayBuffer(100000) self.eviction_queue = [] def compute_state(self, obj_id): """为每个缓存对象计算状态特征""" obj = self.cache[obj_id] return np.array([ obj.size / self.max_size, # 归一化大小 obj.access_count / self.max_count, # 相对频率 (now - obj.last_access) / self.max_interval, # 时间间隔 (now - obj.first_access) / self.max_interval,# 内容年龄 obj.content_type_embedding, # 内容类型编码 self.predicted_next_access(obj_id), # 预测下次访问时间 ]) def select_victim(self): """选择驱逐对象 - 值越小越容易被驱逐""" victim_scores = {} for obj_id in self.cache: state = self.compute_state(obj_id) # 负值表示"预测不会再被访问" victim_scores[obj_id] = self.policy_net.predict(state) # 移除得分最低的(最不需要保留的) return min(victim_scores, key=victim_scores.get) def on_hit(self, obj_id): """缓存命中时的状态更新""" obj = self.cache[obj_id] obj.access_count += 1 obj.last_access = now # 存入 replay buffer 用于训练 self.replay_buffer.add(self.compute_state(obj_id), obj) def train_step(self): """从 replay buffer 中采样训练策略网络""" batch = self.replay_buffer.sample(256) # 奖励 = 命中次数 / 缓存空间占用 # 目标:在有限空间内最大化命中率 self.policy_net.update(batch) ```

    3.3 实际效果与取舍

    Meta 在 CDN 中部署 learned cache 后的关键发现:

    场景 LRU 命中率 Learned Cache 命中率 提升
    视频 CDN (长尾分布) 72% 81% +12.5%
    图像 CDN (强时间局部性) 85% 87% +2.4%
    API 缓存 (突发流量) 54% 63% +16.7%
    数据库缓冲池 78% 79% +1.3%

    关键洞察:访问模式越不可预测,学习策略相比 LRU 的优势越明显。在高度时间局部性的场景下,LRU 本身就是接近最优的策略。

    3.4 轻量级替代:TinyLFU + W-TinyLFU

    对于不愿承担 RL 训练开销的场景,TinyLFU 是一个实用的折中:

    ```python class WTinyLFU: """Window-TinyLFU 缓存方案(Caffeine 缓存库采用)""" def __init__(self, max_size): self.max_size = max_size # Window cache: 使用 SLRU 策略,占 1% 空间 self.window_cache = SLRU(int(max_size * 0.01)) # Main cache: 使用 Segmented LRU,占 99% 空间 self.main_cache = SegmentedLRU(int(max_size * 0.99)) # 频率 Sketch: Count-Min Sketch self.frequency_sketch = CountMinSketch(width=100000, depth=4) # 准入过滤器: Bloom Filter 变体 self.admission_filter = BloomFilter(size=50000) def get(self, key): # 先在 window 中查找 value = self.window_cache.get(key) if value: self.record_access(key) return value # 再在 main 中查找 value = self.main_cache.get(key) if value: self.record_access(key) return value return None def put(self, key, value): # 新条目进入 window cache self.window_cache.put(key, value) # 如果 window 满了,驱逐到 main if self.window_cache.is_full(): evicted_key, evicted_value = self.window_cache.evict() # 比较 window 驱逐者与 main 候选者的频率 if self.should_admit(evicted_key, self.main_cache.admission_candidate()): self.main_cache.put(evicted_key, evicted_value) def should_admit(self, candidate_key, victim_key): """基于频率决定是否允许新条目进入 main""" candidate_freq = self.frequency_sketch.estimate(candidate_key) victim_freq = self.frequency_sketch.estimate(victim_key) return candidate_freq > victim_freq ```

    Caffeine(JVM 生态最快的缓存库)采用 W-TinyLFU,在实际工作负载中显著优于 LRU,且不需要在线学习。


    四、AI 驱动查询优化器

    4.1 传统查询优化器的瓶颈

    传统基于代价的查询优化器 (CBO) 依赖:

    1. 统计信息:直方图、NDV(不同值数量)
    2. 代价模型:I/O 成本 + CPU 成本的加权
    3. 连接顺序枚举:动态规划或贪婪算法
    4. 这些方法的固有缺陷包括:

      • 统计信息过时或不准确
      • 代价模型无法捕捉硬件特性(缓存层级、SIMD 并行度)
      • 连接顺序搜索空间随表数量指数增长

      4.2 Bao:基于贝叶斯优化的学习型优化器

      Bao (Bayesian Optimizer for Analytics) 是近年来最具影响力的学习型查询优化器:

      ```python class BaoOptimizer: """Bao 查询优化器核心逻辑""" def __init__(self): # 树-structured Parzen estimator 用于超参数优化 self.encoder = TreeLSTM() # 查询计划编码为向量 self.latency_predictor = MLP() # 预测各 plan 的延迟 self.exploration_strategy = BayesianExploration() def optimize_query(self, query_ast): # Step 1: 生成候选 plan 集合 candidate_plans = self.generate_candidate_plans(query_ast) # Step 2: 用 TreeLSTM 编码每个 plan plan_embeddings = {} for plan in candidate_plans: plan_embeddings[plan.id] = self.encode_plan(plan) # Step 3: 贝叶斯优化选择最优 plan # 先验:训练数据集中的 (plan, latency) 对 # 后验:在当前工作负载上的预测 predicted_latencies = self.latency_predictor.predict_batch( list(plan_embeddings.values()) ) # Step 4: 探索 vs 利用权衡 # 以一定概率选择一个非最优 plan 来收集训练数据 if self.exploration_strategy.should_explain(): selected_plan = self.exploration_strategy.select(candidate_plans) else: selected_plan = candidate_plans[argmin(predicted_latencies)] # Step 5: 执行并记录实际延迟作为在线训练数据 actual_latency = self.execute_and_measure(selected_plan) self.training_data.append((plan_embeddings[selected_plan.id], actual_latency)) # 定期在后台重新训练预测器 if len(self.training_data) % 100 == 0: self.retrain_predictor() return selected_plan def encode_plan(self, plan): """将查询计划树编码为固定长度向量""" # Tree-LSTM 自底向上编码每个算子 def encode_node(node): if node.is_leaf(): return self.encode_scan(node) child_hs = [encode_node(c) for c in node.children] return node.operator.aggregate(child_hs) return encode_node(plan.root) ```

      4.3 生产部署关键挑战

      Bao 在实际部署中暴露了学习优化器的共性问题:

      1. 泛化性: 训练过的 workload 上表现良好,但面对全新查询模式时可能选择灾难性的计划
      2. 冷启动: 新部署时需要收集足够的训练数据
      3. 可解释性: DBA 难以理解为什么 AI 选择了某个计划(黑盒问题)
      4. 解决方案是混合架构:学习优化器负责选择连接顺序,传统规则保证安全边界(例如不允许全表扫描大表)。


        五、深度强化学习调度器

        5.1 Google Borg 中的调度挑战

        Google 的集群调度器 Borg 每天处理数十万个作业。核心决策:

        • 放置决策: 将作业分配到哪台机器
        • 资源分配: 给作业分配多少 CPU/内存
        • 优先级抢占: 高优先级作业何时抢占低优先级

        传统方法使用基于优先级和约束的规则系统。强化学习调度器的目标是学习一个放置策略,最小化全局资源碎片和作业等待时间。

        5.2 DeepRM 架构

        ```python class DeepRMScheduler: """基于深度强化学习的集群调度器""" def __init__(self, num_resources): # 状态空间:机器可用资源、作业资源需求、等待队列 self.state_dim = num_resources * 2 + JOB_FEATURES # 动作空间:选择将作业分配给哪台机器 self.action_dim = MAX_MACHINES self.policy_network = ActorCritic( state_dim=self.state_dim, action_dim=self.action_dim, hidden_dims=[128, 64, 32] ) def get_state(self): """编码当前集群状态""" cluster_state = [] # 每台机器的可用资源 for machine in self.machines: cluster_state.extend([ machine.available_cpu, machine.available_mem, machine.available_gpu, ]) # 等待队列中作业的特征 for job in self.waiting_queue[:MAX_JOBS]: cluster_state.extend([ job.cpu_request, job.mem_request, job.queue_time, job.priority, ]) return np.array(cluster_state, dtype=np.float32) def compute_reward(self, job, machine): """奖励函数:鼓励紧凑放置 + 减少等待时间""" # 利用率高的小碎片奖励 utilization_reward = machine.forecast_utilization_after(job) # 等待时间惩罚 wait_penalty = -job.queue_time # 负载均衡:避免热点机器 load_balance_reward = -machine.load_variance_increase(job) return utilization_reward + 0.1 * wait_penalty + load_balance_reward def schedule(self, job): state = self.get_state() action_probs = self.policy_network.actor(state) # 资源可行性过滤 feasible_mask = self.get_feasible_machines(job) masked_probs = action_probs * feasible_mask masked_probs /= masked_probs.sum() # 采样机器 machine_idx = np.random.choice(len(masked_probs), p=masked_probs) self.assign_job(job, self.machines[machine_idx]) ```

        5.3 从仿真到生产的鸿沟

        Google 内部的经验表明,强化学习调度器在仿真环境中相比传统调度器提升 15-25% 的资源利用率,但直接部署到生产环境会遇到:

        1. 安全约束: RL 可能做出灾难性决策(如将关键作业放到即将下线的机器上)
        2. 策略更新期间的抖动: 在线学习可能导致调度行为不一致
        3. 多目标优化: 资源利用率、公平性、SLO 满足率之间的权衡难以用单一奖励函数表达
        4. 工业界的务实做法是将 RL 建议作为参考,最终决策仍由规则引擎约束(类似自动驾驶中的 L2 级辅助:AI 建议 + 安全护栏)。


          六、网络拥塞控制的智能化

          6.1 Remy:自动设计的拥塞控制

          MIT 的 Remy 项目使用离线搜索自动生成拥塞控制算法:

          ```python class RemyGenerator: """Remy 的离线自动化拥塞控制设计""" def __init__(self, network_model, traffic_model, design_objectives): self.network = network_model # 网络拓扑模型 self.traffic = traffic_model # 流量模式 self.objectives = design_objectives # 吞吐量、延迟、公平性权重 def generate_cc(self): """搜索最优的拥塞控制策略""" # 状态机:有限的拥塞状态(慢启动、拥塞避免、快速恢复等) # 触发器:ACK 到达、丢包、ECN、定时器超时 # 动作:cwnd 增加/减少/保持 best_strategy = None best_score = -inf # 使用蒙特卡洛树搜索 (MCTS) 搜索状态机 for iteration in range(MAX_ITERATIONS): candidate = self.mcts_search() score = self.evaluate(candidate) if score > best_score: best_score = score best_strategy = candidate return best_strategy def evaluate(self, cc_strategy): """在模拟网络中评估策略""" throughput = self.simulate(cc_strategy, metric='throughput') latency = self.simulate(cc_strategy, metric='latency') fairness = self.simulate(cc_strategy, metric='fairness_score') return (self.objectives.throughput_weight * throughput - self.objectives.latency_weight * latency + self.objectives.fairness_weight * fairness) ```

          Remy 生成的算法在特定网络条件下超越了人类设计的 TCP Cubic 和 BBR,但策略是针对特定网络环境优化的,缺乏泛化性。

          6.2 PCC Vivace:在线学习拥塞控制

          PCC (Performance-oriented Congestion Control) Vivace 使用在线梯度下降优化效用函数:

          ``` 效用函数 U = α * throughput - β * latency_loss - γ * loss_rate 每个 RTT 测量实际 rate 和 utility: 如果 utility 增加 → 继续增大发送速率 如果 utility 减少 → 减小发送速率 通过梯度方向更新速率 ```

          PCC Vivace 的优势在于不需要预先设置拥塞窗口增长规则,完全靠在线学习适应网络条件。


          七、生产部署的工程实践

          7.1 训练数据管道

          AI for Systems 的关键基础设施是训练数据闭环:

          ```python class SystemTrainingPipeline: """系统机器学习组件的训练管道""" def __init__(self): self.feature_store = FeatureStore() self.model_registry = ModelRegistry() self.experiment_tracker = ExperimentTracker() def collect_training_data(self): """在线收集训练样本(对系统性能影响最小化)""" # 1. 在线采样:使用 epsilon-greedy 探索收集数据 # 2. shadow mode: 模型只做预测但不影响决策,记录预测 vs 实际结果 # 3. replay: 离线重放日志生成训练数据 pass def train_model(self, training_data): """在专用训练节点上训练""" # 关键:训练任务不能影响在线服务性能 # 使用 spot/preemptible 训练资源 # 定期 checkpoint 到模型仓库 pass def deploy_model(self, model, rollout_strategy='canary'): """渐进式发布新模型""" # canary: 1% 流量 → 10% → 50% → 100% # 监控关键指标:命中率、延迟、错误率 # 自动回滚机制 pass def monitor_drift(self): """检测模型漂移""" # 输入特征分布变化 # 预测准确率下降 # 系统级指标恶化 # 触发自动再训练 pass ```

          7.2 安全护栏 (Safety Guardrails)

          无论模型多聪明,生产系统都需要安全护栏:

          ```python class GuardedLearnedComponent: """带安全护栏的学习组件""" def __init__(self, learned_model, fallback_policy, max_deviation): self.model = learned_model self.fallback = fallback_policy # 传统启发式策略 self.max_deviation = max_deviation def decide(self, input_data): # 获取学习模型的决策 ml_decision = self.model.predict(input_data) # 获取传统策略的决策 traditional_decision = self.fallback.decide(input_data) # 如果两者偏差太大,使用保守策略 if self.deviation(ml_decision, traditional_decision) > self.max_deviation: return traditional_decision return ml_decision ```

          八、未来展望

          AI for Systems 领域正在从探索走向成熟,几个关键趋势值得关注:

          1. 大语言模型辅助系统调优: 用 LLM 分析系统指标、生成调优建议,甚至自动修改系统参数
          2. 自动特征发现: AutoML for Systems,自动发现系统状态中哪些特征对预测最有价值
          3. 因果推断替代相关性: 从"统计相关"到"因果驱动",确保学习模型的行为可解释、可调试
          4. 联邦学习保护隐私: 跨组织联合训练系统优化模型,无需共享原始数据
          5. 编译器-硬件-算法协同设计: 用 AI 指导整个垂直栈的联合优化
          6. 但核心原则始终不变:学习组件应该是系统的增强,而非替代。最成功的部署案例都采用了"AI 建议 + 传统约束 + 人工审核"的三层架构。


            参考文献

            1. Kraska, T., et al. "The Case for Learned Index Structures." SIGMOD 2018.
            2. Galakatos, A., et al. "Fitting Tree Index Structures to Main Memory Database Platforms." ICDE 2019.
            3. Mao, H., et al. "Learning Cache Replacement Policies." ICML 2019.
            4. Marcus, R., & Papaemmanouil, O. "Bayesian Optimizer for Analytics." VLDB 2019.
            5. Sheng, Y., et al. "Bao: Learning to Steer Query Optimizers." SIGMOD 2021.
            6. Grandl, R., et al. "Altruistic Scheduling in Multi-Resource Clusters." OSDI 2016.
            7. Winstein, K., & Balakrishnan, H. "Remy: Automatic TCP Congestion Control." NSDI 2013.
            8. Dong, M., et al. "PCC Vivace: Online-Learning Congestion Control." NSDI 2018.
            9. Ergin, I., et al. "BaiDao: Baidu's Learned Database." VLDB 2023.
            10. Krishnan, S., et al. "Learning to Optimize Joins." SysML 2018.
点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论
立即
投稿

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部