elrond's Blog

vLLM 推理主线：从一次请求看懂 Attention、KV Cache 与输出生成

Mon, 30 Mar 2026 00:00:00 +0000

1. 引言
2. 一次请求的全景图
3. 请求进入 vLLM 之后的第一段路
4. 模型开始生成时，内部在发生什么
5. 一个 token 如何真正被选出来
6. attention score 与 logits 放在一张图里看
7. vLLM 的性能价值落点
8. 几个高频疑问
9. 总结

vLLM 推理主线：从一次请求看懂 Attention、KV Cache 与输出生成

1. 引言

很多人学大语言模型推理时，会遇到这样一种体验：
每个词都认识，整条链路放在一起就容易散。

Self-Attention 看起来懂一点
Q / K / V 看起来懂一点
KV Cache 也知道和提速有关
LM Head、logits、top-k 这些概念单独拿出来也能解释

到了源码、推理框架或者线上服务，问题就会集中冒出来：

Query 在查什么
KV Cache 里到底放了什么
当前 token 和历史 token 的关系怎么建立
logits 和 attention score 各自在哪一步出现
LM Head 为什么直接连着输出
vLLM 的快，快在模型原理，还是快在工程实现

这篇文章围绕一条线展开：

用户发来一句话之后，vLLM 如何完成输入处理、上下文建模、历史复用，以及下一个 token 的生成。

读完之后，整条链路会形成一个闭环：

用户输入
  ↓
分词
  ↓
向量化
  ↓
多层 Attention 建模
  ↓
建立 KV Cache
  ↓
逐 token 生成
  ↓
LM Head 映射词表
  ↓
logits + 采样
  ↓
输出下一个 token

2. 一次请求的全景图

先把地图立起来，再进入细节。

flowchart TD
    A[用户输入文本] --> B[Tokenizer 分词]
    B --> C[Token IDs]
    C --> D[Embedding + Position Encoding]
    D --> E[Prefill]
    E --> F[多层 Transformer 计算]
    F --> G[历史 K / V 写入 KV Cache]
    G --> H[Decode]
    H --> I[当前位置 Hidden State]
    I --> J[LM Head]
    J --> K[logits]
    K --> L[Top-k / Top-p / Temperature]
    L --> M[选出下一个 token]
    M --> N[拼回上下文]
    N --> H

这张图可以拆成三段：

输入阶段：文本进入模型，形成上下文化语义表示
生成阶段：复用历史缓存，逐 token 往前推进
输出阶段：当前位置的语义结果映射到词表，选出下一个 token

3. 请求进入 vLLM 之后的第一段路

3.1 Tokenizer：文本进入模型的入口

用户输入一句话，例如：

介绍一下北京的旅游景点

这段文本进入 vLLM 后，第一步是 Tokenizer 分词。
模型内部处理的基本单位是 token。

中文场景里，token 的粒度和自然语言里的“词”并不严格一致。
它可能是：

一个汉字
两三个汉字组成的片段
标点
特殊符号

图示如下：

flowchart LR
    A[用户输入文本] --> B[Tokenizer]
    B --> C[token_1]
    B --> D[token_2]
    B --> E[token_3]
    B --> F[token_n]

切分结束后，token 会进一步转成 token id：

"介绍一下北京的旅游景点"
↓
[token_1, token_2, token_3, ..., token_n]
↓
[id_1, id_2, id_3, ..., id_n]

这里有两个直接影响：

输入会被拆成什么粒度
输出阶段每一步有哪些 token 可选

词表在这一步的作用非常直接：
它定义了模型处理语言时使用的基本单位。

3.2 Embedding：token 如何变成向量

token id 还是离散编号。
模型真正计算的对象是向量。

每个 token id 会在 embedding table 中找到自己的向量表示。
同时，还会加上位置信息，告诉模型：

这个 token 在序列中的位置
它和前后 token 的相对关系

图示如下：

flowchart TD
    A[Token IDs] --> B[Embedding Lookup]
    B --> C[Token Vectors]
    C --> D[Position Encoding]
    D --> E[输入表示 X]

到这里，字符串已经变成一串向量序列。
Transformer 的后续计算都基于这组向量展开。

3.3 Q / K / V：注意力计算里的三个角色

对于某个 token 当前的输入表示 x，模型会通过三组参数得到三种投影结果：

Q = x @ W_Q
K = x @ W_K
V = x @ W_V

图示如下：

flowchart LR
    X[某个 token 的输入表示 x] --> Q[Q = xW_Q]
    X --> K[K = xW_K]
    X --> V[V = xW_V]

这三个向量在计算中承担着不同角色：

Q（Query）：当前位置带着什么需求去看上下文
K（Key）：当前位置暴露出哪些特征，方便别的位置与它匹配
V（Value）：当前位置最终提供给别的位置的内容

这个设计带来的好处很清楚：
同一个 token 表示，通过不同投影，可以参与不同职责的计算。

3.4 Self-Attention：上下文中的语义汇聚

Self-Attention 决定了每个 token 如何从上下文中吸收信息。

以当前位置 i 为例，它会做这样一组动作：

拿自己的 Q_i
和上下文所有位置的 K_1 ... K_n 做相关性计算
得到一组 attention score
经过 softmax 变成注意力权重
对上下文中的 V_1 ... V_n 做加权求和
得到当前位置新的表示

图示如下：

flowchart LR
    Q[当前位置 Query] --> S[与上下文所有 Key 计算相关性]
    K1[Key_1] --> S
    K2[Key_2] --> S
    K3[Key_3] --> S
    Kn[Key_n] --> S

    S --> W[注意力权重]
    W --> V1[Value_1]
    W --> V2[Value_2]
    W --> V3[Value_3]
    W --> Vn[Value_n]

    V1 --> O[加权求和]
    V2 --> O
    V3 --> O
    Vn --> O

    O --> R[当前位置新的语义表示]

这一阶段完成的是上下文建模。
当前位置会根据整段上下文，形成更贴合当前语境的表示。

举个非常直观的例子：

苹果今天涨了
我今天买了苹果吃

两个句子中都出现了“苹果”，
经过 attention 之后，它们拿到的上下文信息完全不同，最终表示也会不同。

于是，token 在句子里的含义会被动态塑形。

3.5 Prefill：整段 prompt 的首次计算

当用户 prompt 第一次进入模型时，整段上下文会一起参与计算。
这个阶段叫 prefill。

prefill 的工作内容包括：

prompt 全量进入 Transformer
所有已知 token 计算 Q / K / V
多层 attention 和 FFN 完成整段上下文建模
各层历史 token 的 K / V 写入缓存

图示如下：

flowchart TD
    A[完整 Prompt] --> B[Tokenizer]
    B --> C[Embedding + Position]
    C --> D[多层 Transformer]
    D --> E[计算所有已知 token 的 Q/K/V]
    E --> F[形成上下文化表示]
    E --> G[历史 K/V 写入 KV Cache]

prefill 往往是长上下文场景里成本较高的一段。
因为它需要处理整段 prompt，而这一段计算还没有可复用的历史缓存可拿。

4. 模型开始生成时，内部在发生什么

4.1 Decode：逐 token 推进的生成过程

prefill 结束后，模型进入 decode 阶段。
这时生成开始逐 token 推进。

每一步大致如下：

当前上下文已经存在
新位置计算当前 token 的 Q / K / V
当前 Query 去读取历史 KV Cache
得到当前位置的 hidden state
当前 token 的 K / V 写回缓存
进入输出阶段

图示如下：

flowchart TD
    A[当前上下文] --> B[当前新位置]
    B --> C[计算当前 token 的 Q/K/V]
    C --> D[当前 Query 读取历史 KV Cache]
    D --> E[得到当前位置 hidden state]
    C --> F[当前 K/V 写回 KV Cache]
    E --> G[进入输出阶段]

decode 过程中，模型每一步都只做增量推进。
上下文越长，KV Cache 的价值越明显。

4.2 KV Cache：历史信息如何被复用

KV Cache 存的是历史 token 在各层中的 Key 和 Value。
这样一来，生成新 token 时，历史部分无需重新计算。

图示如下：

flowchart LR
    H1[历史 token 1] --> KV1[K_1 / V_1]
    H2[历史 token 2] --> KV2[K_2 / V_2]
    H3[历史 token 3] --> KV3[K_3 / V_3]
    Hn[历史 token n] --> KVn[K_n / V_n]

    KV1 --> CACHE[KV Cache]
    KV2 --> CACHE
    KV3 --> CACHE
    KVn --> CACHE

    Q[当前 token Query] --> READ[读取历史 K/V]
    CACHE --> READ
    READ --> OUT[当前位置上下文结果]

这套机制决定了生成阶段的性能上限。
历史 token 的 K / V 已经在缓存里，当前步只需要为新增 token 计算对应部分。

这一点和“重新处理整段 prompt”之间的差异，会直接转化为速度差异。

4.3 当前 token 与历史 token 的交互方式

当前 token 与历史 token 的交互可以写成：

attention_scores = Q_t @ K_cache.T
attention_weights = softmax(attention_scores)
context = attention_weights @ V_cache

这里有三个关键点：

Query 来自当前新位置
Key / Value 来自历史缓存
attention score 对应的是“上下文位置”的打分

这一步完成后，当前位置的上下文结果才会形成，随后再进入下一层或进入输出映射。

5. 一个 token 如何真正被选出来

5.1 Hidden State：当前位置的语义结果

当前位置经过多层 Transformer 计算后，会得到一个最终 hidden state。
它是当前位置在当前上下文下的语义浓缩结果。

可以把它理解为：

模型对当前位置的综合判断结果

这个向量已经包含了：

当前上下文的主题
历史 token 的依赖关系
当前生成阶段的语义重心

5.2 LM Head：从语义空间到词表空间

LM Head 对应模型的输出层。
它负责把当前位置的 hidden state 映射到整个词表。

一般写法如下：

logits = h_t @ W_vocab

其中：

h_t：当前位置 hidden state
W_vocab：映射到整个词表空间的参数矩阵
logits：词表维度上的原始分数

图示如下：

flowchart LR
    H[当前位置 hidden state] --> L[LM Head]
    L --> O[词表 logits]

LM Head 接住的是语义空间里的结果，输出的是词表空间里的分数。
这一步把“当前位置想表达什么”翻译成“词表里的哪些 token 分数更高”。

5.3 logits：词表上的原始分数

logits 对应的是整个词表上的原始打分。
假设词表大小是 100000，那么当前位置会得到长度为 100000 的分数向量。

示意如下：

hidden state
   ↓
LM Head
   ↓
[
  token_1: score_1,
  token_2: score_2,
  token_3: score_3,
  ...
  token_n: score_n
]

这些分数表示当前位置输出各个候选 token 的倾向程度。

例如在“介绍一下北京的旅游___”这个上下文下：

“景点”可能分数更高
“历史”可能也有一定分数
标点、终止符等也会参与竞争

5.4 top-k / top-p / temperature：采样规则如何工作

logits 进入采样阶段后，常见的控制参数包括：

temperature：调节分布的平滑程度
top-k：保留分数最高的 k 个候选
top-p：保留累计概率达到阈值 p 的候选集合

流程图如下：

flowchart TD
    A[logits] --> B[Softmax]
    B --> C[Temperature 调整]
    C --> D[Top-k / Top-p 筛选]
    D --> E[采样得到下一个 token]

采样结束后，真正输出的 token 会拼接到当前上下文末尾。
接下来继续进入下一轮 decode。

6. attention score 与 logits 放在一张图里看

这两个概念最容易混。

它们出现的位置不同，打分对象也不同。

attention score

出现在 Self-Attention 阶段
打分对象是历史上下文位置
解决的是当前 token 该关注哪些历史 token

logits

出现在 LM Head 之后
打分对象是整个词表候选 token
解决的是当前位置输出哪个 token

放到一张图里：

flowchart LR
    A[当前 token Query] --> B[历史 Key]
    B --> C[attention score]
    C --> D[历史 Value 加权汇总]
    D --> E[当前位置 hidden state]
    E --> F[LM Head]
    F --> G[logits]
    G --> H[词表采样]

一条更直观的理解路径是：

attention score 负责“看谁”
logits 负责“选谁”

7. vLLM 的性能价值落点

Transformer 的基础计算流程在各类实现里是共通的。
vLLM 的优势主要体现在推理工程层。

7.1 PagedAttention

vLLM 将 KV Cache 拆成固定大小的 block，以分页方式管理显存中的缓存内容。
这让缓存的分配、扩容、回收都更灵活。

图示如下：

逻辑序列位置
   ↓
Block Table
   ↓
KV Blocks
   ↓
物理显存块

在长上下文和多请求并发场景里，这种布局能明显减轻显存碎片问题。

7.2 Block Table 与 KV Block

每个序列会维护一张 block table。
这张表负责描述：

逻辑上的上下文位置
物理上的 KV block 映射关系

这让同一个请求的缓存可以按块组织，随着生成过程动态增长。
多个请求并发时，块的回收与复用也更容易调度。

7.3 Continuous Batching

vLLM 支持连续批处理。
请求可以动态进入批次，也可以动态退出批次。

这对线上服务很关键：

GPU 利用率更高
等待时间更短
不同长度请求能够更自然地混跑

Continuous Batching 和 KV Cache 管理放在一起看，才能看清 vLLM 的吞吐优势从哪里来。

7.4 Prefix Cache

许多请求拥有共同前缀。
聊天场景、带系统提示词的应用、多轮对话场景都很常见。

前缀复用能直接节省 prefill 成本。
这一点在重复模板、重复系统提示、重复历史上下文的业务场景里很有价值。

8. 几个高频疑问

8.1 KV Cache 存的是词表信息吗

KV Cache 存的是当前请求历史 token 的 K / V。
缓存对象属于上下文序列空间。

8.2 Query 没有命中缓存时会发生什么

模型会对当前请求中实际出现的 prompt token 做 prefill，建立对应的 K / V。
计算范围仍然属于当前序列中的 token。

8.3 需要对整个词表计算 K / V 吗

不需要。
词表参与的是 LM Head 和 logits 采样阶段。
K / V 的计算对象属于当前序列里的 token 表示。

8.4 top-k 筛选的是历史 token 吗

top-k 作用在词表采样阶段。
筛选对象是词表候选 token。

8.5 LM Head 和 Self-Attention 的关系怎么理解

Self-Attention 完成上下文建模，形成当前位置的 hidden state。
LM Head 接过这个结果，把它映射到整个词表。

9. 总结

从一次请求的完整生命周期来看，vLLM 中的推理主线可以归纳成下面这张图：

用户输入
  ↓
Tokenizer 分词
  ↓
Embedding + Position Encoding
  ↓
Prefill：整段 prompt 建立上下文表示
  ↓
历史 K/V 写入 KV Cache
  ↓
Decode：逐 token 生成
  ↓
当前位置 hidden state
  ↓
LM Head
  ↓
logits
  ↓
top-k / top-p / temperature
  ↓
下一个 token
  ↓
拼回上下文继续生成

把这条线打通之后，几个核心概念会自然落位：

Q / K / V：注意力计算里的三种角色
Self-Attention：上下文建模的核心机制
Prefill：整段 prompt 的首次计算
KV Cache：历史 K / V 的复用机制
Decode：逐 token 推进的生成过程
LM Head：从 hidden state 映射到整个词表
logits：词表上的原始分数
top-k / top-p：从词表候选中完成采样筛选

站在工程实现角度看，vLLM 的价值集中在：

更高效的 KV Cache 管理
更灵活的显存分页方式
更强的批处理调度能力
更好的在线推理吞吐表现

整条链路清楚之后，再去看 vLLM 的调度器、block manager、paged attention kernel、cache manager，主干会非常清晰。

Vllm：高性能大语言模型推理框架源码解析与最佳实践

Fri, 28 Mar 2025 00:00:00 +0000

vLLM：高性能大语言模型推理框架源码解析与最佳实践

引言
快速上手
2.1. 安装配置
2.2. 基本用法
核心调用流程分析
3.1. 总体调用链路概述
3.2. 核心组件与类层次结构
3.3. 初始化阶段详细流程
3.4. 推理阶段详细流程
3.5. 完整调用链路示例
3.6. 关键调用路径总结
vLLM 关键工作机制
4.1. PagedAttention 机制
4.2. 连续批处理技术
4.3. CUDA 图捕获与加速
4.4. KV 缓存管理
4.5. 内存优化策略
4.6. 并行计算优化
架构与类关系图
5.1. 整体架构概览
5.2. 核心类结构图
5.3. 关键组件交互流程
5.4. 并行与分布式架构
5.5. 内存管理架构
高级 Python 语法应用
6.1. 类型注解与泛型编程
6.2. 动态类解析与反射机制
6.3. 方法委托与魔术方法
6.4. 弱引用与特殊调用模式
6.5. 类方法与工厂模式
6.6. 装饰器高级应用
6.7. 上下文管理与多重上下文
性能优化与最佳实践
7.1. 示例应用：构建高性能API服务
7.2. 性能优化策略
7.3. 部署最佳实践
7.4. 常见问题与解决方案
总结与展望
8.1. 技术创新点总结
8.2. 与其他框架对比
8.3. 未来发展趋势
8.4. 实践建议

附录

A. 关键源码目录结构
B. 参考资源

1. 引言

vLLM 是一个高性能的大语言模型推理框架，通过创新的 PagedAttention 机制和优化的内存管理，实现了高吞吐量和低延迟的文本生成服务。本文将深入分析 vLLM 的源码，重点关注其调用流程、高级 Python 语法应用以及核心工作机制，帮助读者更好地理解和使用这一强大工具。

2. 快速上手

2.1 安装配置

vLLM 支持多种硬件平台，安装方式也略有不同：

# 基本安装（CUDA 支持）
pip install vllm

# 从源码安装
git clone https://github.com/vllm-project/vllm.git
cd vllm
pip install -e .

2.2 基本用法

下面是一个简单的 vLLM 使用示例，展示了如何加载模型并生成文本：

# 导入必要的库
from vllm import LLM, SamplingParams

# 初始化模型
llm = LLM(model="/opt/data/models/Llama-3.2-1B-Instruct")

# 设置生成参数
sampling_params = SamplingParams(
    max_tokens=100,         # 最大生成长度
    temperature=0.7,        # 控制随机性
    top_p=0.9,              # 控制取样范围
    stop=["</s>", "\n\n"]   # 停止标记
)

# 准备输入提示
prompts = [
    "介绍一下北京的旅游景点",
    "写一首关于春天的诗",
]

# 生成文本
outputs = llm.generate(prompts, sampling_params)

# 打印结果
for output in outputs:
    print(f"Prompt: {output.prompt}")
    print(f"Generated: {output.outputs[0].text}")
    print("-" * 50)

3. 核心调用流程分析

3.1 总体调用链路概述

vLLM 框架的调用流程可以分为初始化阶段和推理阶段两个主要部分。整个流程涉及多个核心组件的协同工作，包括 LLM Engine、Worker、Scheduler、ModelRunner 等类。以下是端到端完整调用链路的详细分析：

用户 API 调用
  ↓
LLMEngine 初始化
  ↓ 
模型加载与配置初始化
  ↓
Worker 池初始化
  ↓
请求处理与调度
  ↓
Tokenizer 处理输入
  ↓
ModelRunner 执行推理
  ↓
KV 缓存管理
  ↓
结果获取与后处理
  ↓
返回生成结果

3.2 核心组件与类层次结构

vLLM 的架构由以下核心组件组成：

LLMEngine：整个框架的中枢，负责协调各组件工作
Worker：实际执行模型计算的组件
ModelRunner：负责模型的前向传播计算
Scheduler：请求调度器，决定何时处理哪些请求
SamplingParams：控制文本生成的参数
Request/Response：请求和响应的数据结构
SequenceGroup/Sequence：表示生成序列的数据结构
BlockManager：管理 KV 缓存的物理块
Tokenizer：负责文本标记化

类层次结构图：

LLMEngine
  ├── AsyncEngineDeadlockMonitor
  ├── ModelConfig
  ├── DeviceConfig
  ├── ParallelConfig
  ├── Scheduler
  │    ├── SchedulerConfig
  │    └── SchedulingStrategy
  ├── RequestTracker
  └── WorkerPool
       ├── Worker
       │    ├── ModelRunner
       │    │    ├── CUDAGraph
       │    │    └── AttentionState
       │    └── BlockSpaceManager
       │         └── BlockAllocator
       └── ModelReplica

3.3 初始化阶段详细流程

3.3.1 LLMEngine 初始化

LLMEngine 是整个 vLLM 的核心，负责协调各组件工作。它的初始化流程如下：

def __init__(
    self,
    model_config: ModelConfig,
    scheduler_config: SchedulerConfig,
    device_config: Optional[DeviceConfig] = None,
    parallel_config: Optional[ParallelConfig] = None,
    **kwargs,
) -> None:
    # 1. 设置基本配置
    self.model_config = model_config
    self.scheduler_config = scheduler_config
    if device_config is None:
        device_config = DeviceConfig()
    self.device_config = device_config
    if parallel_config is None:
        parallel_config = ParallelConfig()
    self.parallel_config = parallel_config
    
    # 2. 初始化策略配置
    self.num_nodes = parallel_config.world_size // parallel_config.tensor_parallel_size
    self.distributed_init_method = kwargs.get("distributed_init_method", None)
    self.max_logprobs = kwargs.get("max_logprobs", None)
    
    # 3. 初始化组件
    self.tokenizer = self._init_tokenizer()
    self.schedulers = self._init_schedulers()
    self.request_tracker = RequestTracker()
    
    # 4. 初始化 Worker 池
    self.worker_pool = WorkerPool(
        model_config=model_config,
        parallel_config=parallel_config,
        scheduler_config=scheduler_config,
        device_config=device_config,
        local_rank=parallel_config.local_rank,
    )
    
    # 5. 启动引擎线程
    self._start_background_thread()

初始化过程中的配置详解：

ModelConfig：包含模型相关的配置，如模型路径、类型、量化设置等

class ModelConfig:
    def __init__(
        self,
        model: str,
        tokenizer: Optional[str] = None,
        tokenizer_mode: str = "auto",
        trust_remote_code: bool = False,
        dtype: str = "auto",
        quantization: Optional[str] = None,
        revision: Optional[str] = None,
        code_revision: Optional[str] = None,
        tokenizer_revision: Optional[str] = None,
        # ... 更多参数
    ) -> None:
        # 初始化模型配置

DeviceConfig：包含设备相关配置，如GPU数量、每个GPU的最大内存等

class DeviceConfig:
    def __init__(
        self,
        device_type: Optional[str] = None,
        max_memory_per_gpu: Optional[Dict[int, str]] = None,
        # ... 更多参数
    ) -> None:
        # 初始化设备配置

ParallelConfig：包含并行计算相关配置，如张量并行度、流水线并行度等

class ParallelConfig:
    def __init__(
        self,
        tensor_parallel_size: int = 1,
        pipeline_parallel_size: int = 1,
        # ... 更多参数
    ) -> None:
        # 初始化并行配置

3.3.2 Tokenizer 初始化

Tokenizer 负责将输入文本转换为模型可以处理的 token ID 序列：

def _init_tokenizer(self) -> PreTrainedTokenizer:
    """初始化 tokenizer"""
    # 加载 tokenizer
    tokenizer = get_tokenizer(
        model_name=self.model_config.model,
        tokenizer_name=self.model_config.tokenizer,
        tokenizer_mode=self.model_config.tokenizer_mode,
        trust_remote_code=self.model_config.trust_remote_code,
        revision=self.model_config.tokenizer_revision,
    )
    
    # 配置特殊 token
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    # 检查并处理特殊 token ID
    if (self.model_config.chat_template is not None and
            self.model_config.chat_template != "none"):
        tokenizer.chat_template = self.model_config.chat_template
        
    return tokenizer

3.3.3 Scheduler 初始化

Scheduler 负责调度请求，决定哪些请求应该被处理以及何时处理：

def _init_schedulers(self) -> List[Scheduler]:
    """初始化调度器"""
    # 创建多个调度器实例，每个对应一个模型副本
    schedulers = []
    for rank in range(self.num_nodes):
        # 创建调度器实例
        scheduler = Scheduler(
            # 配置调度器
            scheduler_config=self.scheduler_config,
            max_num_batched_tokens=self.worker_pool.max_num_batched_tokens,
            max_num_seqs=self.worker_pool.max_num_seqs,
            device=f"cuda:{self.parallel_config.local_rank}" if torch.cuda.is_available() else "cpu",
        )
        schedulers.append(scheduler)
    return schedulers

3.3.4 Worker 池初始化

Worker 池负责管理多个 Worker 实例，处理并发请求：

class WorkerPool:
    def __init__(
        self,
        model_config: ModelConfig,
        parallel_config: ParallelConfig,
        scheduler_config: SchedulerConfig,
        device_config: DeviceConfig,
        local_rank: int,
    ) -> None:
        # 1. 设置基本配置
        self.model_config = model_config
        self.parallel_config = parallel_config
        self.scheduler_config = scheduler_config
        self.device_config = device_config
        self.local_rank = local_rank
        
        # 2. 初始化其他参数
        self.max_num_seqs = scheduler_config.max_num_seqs
        self.max_num_batched_tokens = scheduler_config.max_num_batched_tokens
        
        # 3. 初始化 Worker
        self.workers: List[Worker] = []
        self.model_replicas = []
        
        # 4. 加载模型并创建 Worker
        self._initialize_workers()

Worker 初始化时会加载模型并准备执行环境：

def _initialize_workers(self) -> None:
    """初始化 Worker"""
    # 创建模型副本
    self.model_replicas = self._initialize_model_replicas()
    
    # 为每个模型副本创建一个 Worker
    for i, model_replica in enumerate(self.model_replicas):
        worker = Worker(
            model_config=self.model_config,
            model_replica=model_replica,
            # ... 其他参数
        )
        self.workers.append(worker)

3.4 推理阶段详细流程

3.4.1 请求接收与处理

用户通过 API 发送请求，LLMEngine 接收并处理请求：

def add_request(
    self,
    request_id: str,
    prompt: Optional[str],
    sampling_params: SamplingParams,
    prompt_token_ids: Optional[List[int]] = None,
    arrival_time: Optional[float] = None,
) -> Optional[int]:
    """添加一个新请求"""
    # 1. 记录请求到达时间
    if arrival_time is None:
        arrival_time = time.time()
    
    # 2. 处理输入
    if prompt is None and prompt_token_ids is None:
        raise ValueError("Either prompt or prompt_token_ids must be provided")
    
    # 3. 如果需要，使用 tokenizer 处理输入文本
    if prompt_token_ids is None:
        assert prompt is not None
        prompt_token_ids = self.tokenizer.encode(prompt)
    
    # 4. 创建请求并添加到队列
    request = Request(
        request_id=request_id,
        prompt=prompt,
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
        arrival_time=arrival_time,
    )
    
    # 5. 将请求添加到请求跟踪器
    self.request_tracker.add_request(request)
    
    # 6. 将请求添加到调度器
    self.schedulers[0].add_request(request)
    
    return len(prompt_token_ids)

3.4.2 请求调度

Scheduler 负责决定何时处理哪些请求：

def schedule(self) -> List[SequenceGroupMetadata]:
    """调度请求执行"""
    # 1. 获取系统状态
    now = time.time()
    
    # 2. 从等待队列和运行中队列选择可处理的序列组
    seq_groups = self._get_schedulable_groups()
    
    # 3. 根据调度策略对序列组排序
    seq_groups = self._sort_by_strategy(seq_groups)
    
    # 4. 根据资源约束选择可处理的序列组
    scheduled_groups = self._select_best_fitting_groups(seq_groups)
    
    # 5. 生成调度元数据
    metadata = self._create_schedule_metadata(scheduled_groups)
    
    return metadata

3.4.3 模型执行

ModelRunner 负责执行模型的前向传播：

def execute_model(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    kv_caches: List[torch.Tensor],
    block_tables: torch.Tensor,
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
    """执行模型前向传播"""
    # 1. 准备输入
    batch_size = input_ids.shape[0]
    
    # 2. 如果已启用 CUDA 图捕获并且批大小匹配，使用 CUDA 图执行
    if (self.model_captured and batch_size in self.graph_runners and
            not self.disable_graphs):
        # 使用预先捕获的 CUDA 图执行模型
        return self._execute_with_graph(batch_size, input_ids, positions,
                                       kv_caches, block_tables)
    else:
        # 使用标准方式执行模型
        return self._forward_helper(input_ids, positions, kv_caches,
                                    block_tables)

3.4.4 KV 缓存管理

BlockManager 负责管理 KV 缓存的物理块：

def allocate(self, seq_id: int, num_tokens: int) -> List[int]:
    """为序列分配 KV 缓存块"""
    # 1. 计算需要的块数
    num_blocks = (num_tokens + self.block_size - 1) // self.block_size
    
    # 2. 分配物理块
    block_indices = []
    for _ in range(num_blocks):
        block_idx = self.block_allocator.allocate()
        block_indices.append(block_idx)
    
    # 3. 更新映射表
    self.block_tables[seq_id] = block_indices
    
    return block_indices

3.4.5 结果处理与返回

LLMEngine 处理模型输出并返回生成结果：

def _process_model_outputs(
    self,
    seq_group: SequenceGroup,
    outputs: ModelOutput,
    metadata: SequenceGroupMetadata,
) -> None:
    """处理模型输出"""
    # 1. 处理 logits 和采样结果
    logprobs = outputs.logprobs
    next_tokens = outputs.next_tokens
    
    # 2. 为每个序列更新状态
    for seq_id, seq in seq_group.seqs.items():
        # 获取对应的 token 和概率
        next_token = next_tokens[seq_id]
        logprob = logprobs[seq_id] if logprobs is not None else None
        
        # 更新序列状态
        seq.append_token_id(next_token, logprob)
        
        # 检查是否达到终止条件
        if self._is_finished(seq, next_token):
            seq.status = SequenceStatus.FINISHED

3.5 完整调用链路示例

以下是从用户请求到生成文本的完整调用链路示例：

用户 API 调用：

from vllm import LLM, SamplingParams
   
# 设置采样参数
sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=100
)
   
# 初始化 LLM
llm = LLM(model="gpt-3.5-turbo")
   
# 发送请求
outputs = llm.generate("Tell me a joke", sampling_params=sampling_params)
   
# 获取结果
print(outputs[0].outputs[0].text)

LLM 初始化：

# LLM 类内部初始化 LLMEngine
self.engine = LLMEngine(
    model=model,
    tokenizer=tokenizer,
    # 其他参数...
)

请求处理：

# LLM.generate 方法内部调用 LLMEngine.add_request
request_id = self.engine.add_request(
    prompt=prompt,
    sampling_params=sampling_params
)

请求调度与执行：

# LLMEngine 内部的调度循环
while not engine_stopped:
    # 1. 调度请求
    scheduled_groups = self.scheduler.schedule()
       
    # 2. 将调度的请求发送给 Worker
    for worker, metadata in zip(self.workers, scheduled_groups):
        worker.execute_model(metadata)
       
    # 3. 处理模型输出
    for output in outputs:
        self._process_model_outputs(output)
       
    # 4. 更新请求状态
    self.request_tracker.update_status()

结果返回：

# LLMEngine 将结果返回给 LLM
finished_requests = self.request_tracker.get_finished()
   
# LLM 将结果返回给用户
return [r.get_result() for r in finished_requests]

3.6 关键调用路径总结

vLLM 的核心调用流程可以总结为以下几个关键路径：

初始化路径：

LLM.__init__ → LLMEngine.__init__ → 
[ModelConfig, DeviceConfig, ParallelConfig 初始化] → 
Tokenizer 初始化 → Worker 池初始化 → 
[模型加载, Scheduler 初始化] → 引擎线程启动

请求处理路径：

LLM.generate → LLMEngine.add_request → 
Tokenizer.encode → RequestTracker.add_request → 
Scheduler.add_request → Scheduler.schedule →
Worker.execute_model → ModelRunner.forward → 
[BlockManager 管理 KV 缓存] →
LLMEngine._process_model_outputs → 
LLM 返回结果

并行执行路径：

多个 Worker 并行执行 →
每个 Worker 内部使用 CUDA 图加速 →
使用张量并行和流水线并行提高吞吐量

通过这种模块化的设计，vLLM 能够高效地处理大量并发请求，同时优化内存使用和计算资源分配，为大型语言模型提供高性能的推理服务。

4. vLLM 关键工作机制

vLLM 的高性能源于几项关键技术创新。这些机制共同解决了大语言模型推理中的主要瓶颈问题：内存效率、计算速度和资源利用率。以下是这些核心机制的详细分析，即使没有深厚的技术背景，也能理解它们解决的问题和带来的好处。

5.1 PagedAttention 机制

PagedAttention 是 vLLM 的核心创新，通过分页管理 KV 缓存提高内存效率。这一机制受操作系统中的虚拟内存管理启发，解决了大语言模型推理中的关键瓶颈问题。

KV 缓存的挑战

在标准的 Transformer 解码器中，每生成一个新 token，都需要计算注意力，这涉及与之前所有 token 的 key 和 value 的交互。随着序列长度增加，这些 key 和 value 缓存（即 KV 缓存）占用的内存也线性增长：

# 传统 KV 缓存的内存占用
memory_per_seq = seq_len * num_layers * 2 * hidden_dim * dtype_size

对于长序列和批处理请求，内存消耗非常显著。例如，一个批次中有多个不同长度的序列，传统方法需要为每个序列分配足够大的连续内存块，导致内存碎片和浪费。

PagedAttention 原理

PagedAttention 的核心思想是将 KV 缓存组织成固定大小的物理块，并使用块表映射逻辑位置到物理块：

class PagedAttention:
    def __init__(self, block_size: int):
        self.block_size = block_size
        self.block_tables = {}  # 序列到物理块的映射
        self.physical_blocks = []  # 实际存储块
        
    def allocate_block(self, seq_id: int):
        """为序列分配新块"""
        if seq_id not in self.block_tables:
            self.block_tables[seq_id] = []
            
        # 分配新物理块
        block_id = self._find_free_block()
        self.block_tables[seq_id].append(block_id)
        
    def get_kv_cache(self, seq_id: int, positions: torch.Tensor):
        """获取缓存内容"""
        if seq_id not in self.block_tables:
            return None
            
        # 计算物理块索引和块内位置
        block_indices, block_offsets = self._compute_indices(positions)
        physical_blocks = [self.block_tables[seq_id][idx] for idx in block_indices]
        
        # 获取缓存内容
        return self._gather_blocks(physical_blocks, block_offsets)

工作流程：

内存分块：KV 缓存被划分为固定大小的块（如每块 16 个 token）
动态分配：每个序列动态分配所需的块，而不是预先分配
虚拟寻址：使用块表将逻辑位置映射到物理块
资源复用：当序列完成时，其块可以被释放并重用

技术细节

实际实现中，PagedAttention 包含以下关键组件：

BlockAllocator：管理物理块的分配和释放

class BlockAllocator:
    def __init__(self, num_blocks: int):
        self.num_blocks = num_blocks
        self.free_blocks = set(range(num_blocks))
        self.used_blocks = set()
        
    def allocate(self) -> int:
        """分配一个空闲块"""
        if not self.free_blocks:
            raise RuntimeError("No free blocks available")
        
        block_id = next(iter(self.free_blocks))
        self.free_blocks.remove(block_id)
        self.used_blocks.add(block_id)
        return block_id
        
    def free(self, block_id: int) -> None:
        """释放一个块"""
        self.used_blocks.remove(block_id)
        self.free_blocks.add(block_id)

BlockTable：维护序列到物理块的映射

class BlockTable:
    def __init__(self, block_size: int):
        self.block_size = block_size
        self.tables = {}  # seq_id -> List[block_id]
        
    def add_block(self, seq_id: int, block_id: int) -> None:
        """为序列添加一个块"""
        if seq_id not in self.tables:
            self.tables[seq_id] = []
        self.tables[seq_id].append(block_id)
        
    def get_physical_blocks(self, seq_id: int, positions: torch.Tensor) -> Tuple[List[int], torch.Tensor]:
        """获取物理块 ID 和偏移量"""
        block_ids = self.tables[seq_id]
        block_indices = positions // self.block_size
        block_offsets = positions % self.block_size
        physical_blocks = [block_ids[idx.item()] for idx in block_indices]
        return physical_blocks, block_offsets

注意力计算核心：高效实现分块注意力计算

# CUDA 核心伪代码
@cuda.kernel
def paged_attention_kernel(
    q: Tensor,             # [batch_size, num_heads, head_dim]
    k_cache: Tensor,       # [num_blocks, block_size, num_heads, head_dim]
    v_cache: Tensor,       # [num_blocks, block_size, num_heads, head_dim]
    block_tables: Tensor,  # [batch_size, max_blocks]
    output: Tensor         # [batch_size, num_heads, head_dim]
):
    # 获取线程索引
    batch_idx = cuda.blockIdx.x
    head_idx = cuda.blockIdx.y
    
    # 获取该序列的块表
    seq_blocks = block_tables[batch_idx]
    
    # 本地计算缓冲区
    local_k = shared_memory[...]
    local_v = shared_memory[...]
    
    # 加载查询向量
    query = q[batch_idx, head_idx]
    
    # 对每个块执行注意力计算
    for i in range(len(seq_blocks)):
        block_id = seq_blocks[i]
        
        # 加载 KV 缓存到共享内存
        local_k.copy_from(k_cache[block_id])
        local_v.copy_from(v_cache[block_id])
        
        # 计算注意力分数
        scores = compute_attention(query, local_k)
        
        # 应用注意力权重
        output[batch_idx, head_idx] += apply_attention(scores, local_v)

性能优势

PagedAttention 带来了多项显著优势：

内存效率：通过分块和动态分配，显著减少内存碎片
- 传统方法：为每个序列分配最大长度的连续内存
- PagedAttention：按需分配块，高效共享物理内存
批处理吞吐量：支持更多并发请求
- 实验表明，相同硬件上，vLLM 可以处理比其他框架多 2-4 倍的并发请求
长文本支持：优雅处理长序列生成
- 只需添加新块，无需重新分配或复制整个缓存
内存利用率：内存使用量与实际 token 数成正比，而非预分配
- 典型场景中，与传统方法相比，内存使用可减少 50-70%

实际应用案例

在 vLLM 中，PagedAttention 的应用可以在多方面观察到：

动态批处理：高效支持连续批处理（Continuous Batching）

# 在每次前向传播中添加新序列
def step(self):
    # 调度新序列和正在进行的序列
    scheduled_seq_groups = self.scheduler.schedule()
    if not scheduled_seq_groups:
        return
       
    # 准备批处理输入
    batch = self._prepare_batch(scheduled_seq_groups)
       
    # 为新序列分配 KV 缓存块
    for seq_id in batch.new_seq_ids:
        self.block_manager.allocate_blocks(seq_id, batch.prompt_lens[seq_id])
           
    # 执行模型前向传播
    outputs = self.model_runner.forward(batch.input_ids, batch.block_tables, ...)

高效内存管理：动态分配和回收物理块

def free_finished_sequences(self):
    """释放已完成序列的内存"""
    for seq_id in self.finished_sequences:
        # 获取分配的块
        blocks = self.block_tables.get_blocks(seq_id)
           
        # 释放块
        for block_id in blocks:
            self.block_allocator.free(block_id)
               
        # 移除块表条目
        self.block_tables.remove(seq_id)

缓存共享：支持 fork 和复制操作

def fork_sequence(self, src_seq_id: int, dst_seq_id: int, fork_pos: int):
    """分叉序列，复用 KV 缓存"""
    # 获取源序列的块表
    src_blocks = self.block_tables.get_blocks(src_seq_id)
       
    # 计算 fork 位置对应的块
    fork_block_idx = fork_pos // self.block_size
       
    # 共享前 fork_block_idx 个块
    shared_blocks = src_blocks[:fork_block_idx]
       
    # 为新序列创建块表，共享已有块
    self.block_tables.create(dst_seq_id, shared_blocks)
       
    # 从 fork 位置开始分配新块
    return dst_seq_id

总的来说，PagedAttention 是 vLLM 的关键创新，它解决了大语言模型推理中的内存效率问题，使得服务能够以更低的成本处理更多并发请求，同时支持更长的文本生成。这种设计灵感来自操作系统的虚拟内存管理，证明了经典计算机科学原理在现代 AI 系统中的价值。

5.2 连续批处理技术

连续批处理（Continuous Batching）是 vLLM 中与 PagedAttention 配合使用的关键技术，它解决了传统批处理在 LLM 推理中的效率问题。

传统批处理的局限性

传统的批处理策略（Static Batching）存在几个主要问题：

请求同步问题：所有序列必须同时启动和结束

# 传统批处理示例
def static_batch_inference(model, prompts):
    # 编码所有输入
    input_ids = tokenizer(prompts, padding=True, return_tensors="pt")
       
    # 所有序列同时开始生成
    outputs = model.generate(
        input_ids,
        max_length=max_length,
        num_return_sequences=1,
    )
       
    # 所有序列必须等待最长的序列完成
    return tokenizer.batch_decode(outputs)

计算资源浪费：
- 长序列会使短序列等待
- GPU 利用率低下
- 每批次间存在空闲时间
延迟问题：
- 新请求必须等待当前批次完成
- 平均等待时间长

连续批处理原理

连续批处理允许动态添加和移除序列，无需等待整个批次完成：

class ContinuousBatcher:
    def __init__(self):
        self.active_sequences = {}  # 活跃序列集合
        self.waiting_sequences = Queue()  # 等待队列
        
    def add_sequence(self, seq_id, prompt):
        """添加新序列到等待队列"""
        self.waiting_sequences.put((seq_id, prompt))
        
    def step(self, model):
        """执行一步批处理"""
        # 1. 调度可执行序列
        batch = self._schedule_batch()
        
        # 2. 执行模型推理
        outputs = model.forward(batch.input_ids)
        
        # 3. 处理结果
        for seq_id, output in zip(batch.seq_ids, outputs):
            self._process_output(seq_id, output)
            
        # 4. 移除已完成序列
        self._cleanup_finished_sequences()
        
    def _schedule_batch(self):
        """调度当前步骤的批次"""
        batch = Batch()
        
        # 添加活跃序列
        for seq_id, seq in self.active_sequences.items():
            if not seq.is_finished:
                batch.add(seq_id, seq.get_next_input())
                
        # 尝试添加新序列（如果有空间）
        while len(batch) < self.max_batch_size and not self.waiting_sequences.empty():
            seq_id, prompt = self.waiting_sequences.get()
            self.active_sequences[seq_id] = Sequence(prompt)
            batch.add(seq_id, self.active_sequences[seq_id].get_next_input())
            
        return batch

工作流程：

动态加入：新请求可以在任何时间点加入批处理
并行处理：不同长度和进度的序列可以同时处理
动态退出：完成的序列立即释放资源
资源调度：根据当前负载动态调整批大小

vLLM 的实现细节

在 vLLM 中，连续批处理由 Scheduler 和 Engine 类协同实现：

请求调度器：管理请求队列和调度决策

class Scheduler:
    def __init__(self, max_num_seqs: int, max_num_tokens: int):
        self.max_num_seqs = max_num_seqs  # 最大序列数
        self.max_num_tokens = max_num_tokens  # 最大 token 数
        self.waiting = []  # 等待队列
        self.running = {}  # 运行中的序列
        self.swapped = {}  # 被换出的序列
        
    def schedule(self) -> List[SequenceGroup]:
        """调度下一批序列"""
        scheduled = []
        
        # 优先调度运行中的序列
        running_ready = [seq for seq in self.running.values() if seq.is_ready()]
        scheduled.extend(running_ready)
        
        # 如果有剩余资源，从等待队列调度新序列
        remaining_capacity = self._get_remaining_capacity(scheduled)
        if remaining_capacity.can_schedule():
            new_seqs = self._schedule_from_waiting(remaining_capacity)
            scheduled.extend(new_seqs)
            
        return scheduled
        
    def add_request(self, request):
        """添加新请求到等待队列"""
        self.waiting.append(request)
        
    def abort_request(self, request_id):
        """中止请求"""
        # 从所有队列中移除
        # ...
        
    def _schedule_from_waiting(self, capacity):
        """从等待队列调度序列"""
        scheduled = []
        remaining_seqs = capacity.max_seqs
        remaining_tokens = capacity.max_tokens
        
        # 使用优先级排序等待队列
        self.waiting.sort(key=lambda x: x.priority, reverse=True)
        
        for seq in self.waiting:
            # 检查是否有足够资源
            if remaining_seqs > 0 and remaining_tokens >= seq.token_count:
                scheduled.append(seq)
                self.waiting.remove(seq)
                self.running[seq.id] = seq
                
                # 更新剩余容量
                remaining_seqs -= 1
                remaining_tokens -= seq.token_count
                
        return scheduled

引擎核心：执行实际批处理

class Engine:
    def __init__(self, model, scheduler):
        self.model = model
        self.scheduler = scheduler
        self.running_sequences = {}
        
    def add_request(self, request):
        """添加新请求"""
        self.scheduler.add_request(request)
        
    def step(self):
        """执行一步批处理"""
        # 1. 调度序列
        scheduled = self.scheduler.schedule()
        if not scheduled:
            return None
            
        # 2. 准备批处理输入
        batch = self._prepare_batch(scheduled)
        
        # 3. 执行模型推理
        outputs = self.model.forward(batch.input_ids, batch.attention_mask)
        
        # 4. 处理输出
        self._process_outputs(outputs, batch)
        
        # 5. 返回生成的 token
        return self._get_outputs()
        
    def _prepare_batch(self, scheduled):
        """准备批处理输入"""
        batch = Batch()
        
        for seq in scheduled:
            if seq.id in self.running_sequences:
                # 继续生成
                input_id = self.running_sequences[seq.id].last_token
            else:
                # 新序列，处理提示
                input_ids = self._tokenize(seq.prompt)
                self.running_sequences[seq.id] = RunningSequence(seq.id, input_ids)
                
            batch.add(seq)
            
        return batch
        
    def _process_outputs(self, outputs, batch):
        """处理模型输出"""
        for i, seq_id in enumerate(batch.seq_ids):
            # 获取下一个 token
            next_token = outputs.logits[i].argmax(-1).item()
            
            # 更新序列
            self.running_sequences[seq_id].append_token(next_token)
            
            # 检查是否完成
            if self._is_finished(seq_id, next_token):
                seq = self.running_sequences.pop(seq_id)
                self.scheduler.finish_sequence(seq_id)

优化技术

vLLM 在连续批处理基础上还引入了多项优化：

迭代式调度：每轮迭代动态调整批大小，最大化 GPU 利用率

def iterative_schedule(self):
    """迭代调度以最大化 GPU 利用率"""
    # 初始化批次
    batch = Batch()
    
    # 迭代添加序列，直到达到资源限制
    while True:
        # 尝试添加下一个序列
        next_seq = self._get_next_candidate()
        if next_seq is None:
            break
            
        # 计算添加后的资源使用
        new_resource_usage = self._calculate_resource_usage(batch, next_seq)
        
        # 如果超出资源限制，停止添加
        if new_resource_usage > self.max_resource_limit:
            break
            
        # 添加到批次
        batch.add(next_seq)
    
    return batch

预取优化：提前加载下一批次数据

def prefetch_next_batch(self):
    """预取下一批次数据"""
    # 预测下一步可能调度的序列
    predicted_next = self.scheduler.peek_next_batch()
    
    # 异步预加载数据
    asyncio.create_task(self._async_prefetch(predicted_next))
    
def _async_prefetch(self, sequences):
    """异步预加载数据"""
    for seq in sequences:
        if seq.id not in self.prefetched_data:
            # 预处理输入
            processed_data = self._preprocess(seq)
            
            # 存储预处理结果
            self.prefetched_data[seq.id] = processed_data

优先级调度：基于等待时间和优先级的混合调度策略

def priority_schedule(self):
    """基于优先级的调度"""
    # 计算每个请求的综合优先级分数
    for req in self.waiting:
        # 基础优先级
        priority = req.priority
        
        # 等待时间因子
        wait_time = time.time() - req.arrival_time
        wait_factor = min(1.0, wait_time / self.max_wait_time)
        
        # 综合分数 (混合优先级和等待时间)
        req.score = priority * 0.7 + wait_factor * 0.3
    
    # 按分数排序
    self.waiting.sort(key=lambda x: x.score, reverse=True)
    
    # 调度顶部序列
    return self.waiting[:self.max_batch_size]

性能优势

连续批处理带来的性能优势包括：

吞吐量提升：
- 减少 GPU 空闲时间
- 在相同硬件上支持更多并发请求
响应延迟降低：
- 新请求不必等待整个批次完成
- 短序列可以更快完成，不受长序列影响
资源利用率提高：
- GPU 利用率从传统的 30-40% 提升至 80-90%
- 内存使用更加高效，随时释放完成序列资源
服务质量改进：
- 支持请求优先级
- 更公平的资源分配

实际应用案例

在 vLLM 中，连续批处理结合 PagedAttention，形成了高效的推理引擎：

# vLLM 引擎主循环
def _run_engine(self):
    """vLLM 引擎主循环"""
    while not self._stop_event.is_set():
        # 调度序列
        scheduled_seq_groups = self.scheduler.schedule()
        if not scheduled_seq_groups:
            time.sleep(0.001)  # 避免空转
            continue
            
        # 准备批次
        batch = self._prepare_batch(scheduled_seq_groups)
        
        # 执行模型前向传播
        outputs = self.model_runner.execute_model(
            batch.input_ids,
            batch.positions,
            batch.kv_caches,
            batch.block_tables,
        )
        
        # 后处理结果
        self._process_model_outputs(outputs, batch)
        
        # 检查完成的序列
        self._check_completion(batch.seq_groups)

连续批处理允许 vLLM 在处理大量并发请求时保持高吞吐量和低延迟，这对于构建可扩展的 LLM 服务至关重要。结合 PagedAttention 的内存优化，vLLM 能够高效地管理序列生命周期和计算资源，为大规模 LLM 服务提供关键支持。

5.3 CUDA 图捕获与加速

vLLM 使用 CUDA 图（CUDA Graph）技术显著加速模型推理过程。这种技术通过预先记录和优化 GPU 操作序列，减少 CPU-GPU 交互开销，提高吞吐量和降低延迟。

CUDA 图基本原理

CUDA 图是 NVIDIA GPU 编程中的一项高级功能，允许将一系列 GPU 操作（包括内核启动、内存复制等）预先捕获为一个图，然后重复执行：

# 传统 CUDA 执行方式
for i in range(iterations):
    # 每次迭代都需要 CPU-GPU 通信和调度开销
    kernel1<<<grid, block>>>(args1);  # CPU 调度 GPU 执行
    kernel2<<<grid, block>>>(args2);  # CPU 调度 GPU 执行

# CUDA 图执行方式
# 1. 捕获阶段（只执行一次）
cudaStreamBeginCapture(stream);
kernel1<<<grid, block, 0, stream>>>(args1);
kernel2<<<grid, block, 0, stream>>>(args2);
cudaGraph_t graph;
cudaStreamEndCapture(stream, &graph);

# 2. 实例化图（只执行一次）
cudaGraphExec_t graphExec;
cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0);

# 3. 重复执行（高效）
for i in range(iterations):
    # 无需 CPU 干预，极低开销
    cudaGraphLaunch(graphExec, stream);

CUDA 图的主要优势：

减少 CPU 开销：降低 kernel 调度和参数传递的延迟
优化 GPU 执行：允许 CUDA 运行时预先优化整个计算图
更好的并行性：提高 GPU 流水线效率
确定性执行：一致的执行模式和时间

vLLM 中的 CUDA 图实现

vLLM 在 ModelRunner 类中实现了 CUDA 图捕获机制，主要通过 capture_model 方法完成：

@torch.inference_mode()
def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:
    """捕获模型执行为 CUDA 图以加速推理"""
    # 1. 准备捕获环境
    self._check_if_cuda_graphs_supported()
    # 保存当前内存状态，在捕获后释放临时内存
    mempool = torch.cuda.graph_pool_handle()
    self.graph_memory_pool = mempool
    self.dummy_inputs = {}
    self.graph_runners = {}
    
    # 2. 获取批处理大小
    batch_sizes = self._get_graph_batch_sizes()
    
    # 3. 为每个批处理大小捕获图
    for batch_size in batch_sizes:
        # 准备虚拟模型输入（重要！）
        dummy_inputs = self._prepare_dummy_inputs(batch_size)
        self.dummy_inputs[batch_size] = dummy_inputs
        
        # 捕获模型前向传播
        with self.attn_state.graph_capture(batch_size), \
             cuda_graph_context(self.device) as context:
            # 执行捕获
            outputs = self._forward_helper(**dummy_inputs)
            
        # 创建图运行器
        self.graph_runners[batch_size] = CUDAGraphRunner(
            context.graph, dummy_inputs, outputs
        )
    
    # 4. 设置已捕获标志
    self.model_captured = True

这个方法执行以下核心步骤：

环境检查与设置：验证硬件和软件支持 CUDA 图，并准备内存池
批处理尺寸确定：计算需要支持的批处理大小
准备虚拟输入：为每个批处理大小创建固定形状的输入张量
捕获过程：在特定环境下执行一次模型前向传播，捕获所有 CUDA 操作
图实例化：将捕获的图转换为可执行实例，存储在 graph_runners 中

关键组件：CUDAGraphRunner

CUDAGraphRunner 类封装了 CUDA 图的管理和执行逻辑：

class CUDAGraphRunner:
    """管理 CUDA 图的捕获和执行"""
    
    def __init__(self, graph, static_inputs, static_outputs):
        # 存储图和静态内存引用
        self.graph = graph
        self.static_inputs = static_inputs
        self.static_outputs = static_outputs
        
    def capture(self, **inputs):
        """执行并捕获计算图"""
        # 使用指定输入执行一次前向传播，让 CUDA 记录操作序列
        with torch.cuda.graph(self.graph):
            # 更新静态输入
            for name, tensor in inputs.items():
                self.static_inputs[name].copy_(tensor)
                
            # 捕获前向传播
            self.static_outputs = self.forward_fn(**self.static_inputs)
            
        return self
        
    def forward(self, **inputs):
        """使用CUDA图执行前向传播"""
        # 更新输入数据
        for name, tensor in inputs.items():
            self.static_inputs[name].copy_(tensor)
            
        # 执行图 - 高效！
        self.graph.replay()
        
        return self.static_outputs

这个类负责两个主要任务：

图捕获：在初始化时记录计算图
高效执行：在推理阶段重播图，避免 CPU 调度开销

虚拟输入的重要性

CUDA 图要求输入张量的形状和内存地址在每次执行时保持不变。因此，vLLM 使用”虚拟输入”策略：

def _prepare_dummy_inputs(self, batch_size: int) -> Dict[str, torch.Tensor]:
    """准备用于图捕获的虚拟输入"""
    # 1. 创建固定形状的输入张量
    input_ids = torch.zeros(
        batch_size, dtype=torch.long, device=self.device
    )
    
    # 2. 分配位置 ID
    positions = torch.zeros(
        batch_size, dtype=torch.long, device=self.device
    )
    
    # 3. 创建注意力掩码
    attention_mask = torch.ones(
        batch_size, self.max_context_len, dtype=torch.bool, device=self.device
    )
    
    # 4. 准备 KV 缓存和块表
    # 此处省略复杂的 KV 缓存准备逻辑...
    
    return {
        "input_ids": input_ids,
        "positions": positions,
        "attention_mask": attention_mask,
        "kv_caches": kv_caches,
        "block_tables": block_tables,
    }

在实际执行时，模型会先将真实输入复制到这些预分配的虚拟张量中，然后执行图：

def execute_model(self, input_ids, positions, kv_caches, block_tables):
    """执行模型前向传播"""
    batch_size = input_ids.shape[0]
    
    # 如果已捕获且批处理大小匹配，使用 CUDA 图
    if self.model_captured and batch_size in self.graph_runners:
        # 准备输入
        inputs = {
            "input_ids": input_ids,
            "positions": positions,
            "kv_caches": kv_caches,
            "block_tables": block_tables,
        }
        
        # 使用图运行器 - 高性能路径
        return self.graph_runners[batch_size].forward(**inputs)
        
    else:
        # 标准执行路径
        return self._forward_helper(
            input_ids, positions, kv_caches, block_tables
        )

技术挑战与解决方案

CUDA 图实现面临几个技术挑战：

内存地址固定要求：
- 挑战：CUDA 图要求输入/输出张量地址不变
- 解决方案：使用预分配的静态内存，通过 copy_ 更新内容
动态形状处理：
- 挑战：不同批处理大小需要不同图
- 解决方案：为常见批处理大小预捕获多个图
内存管理：
- 挑战：图捕获可能导致内存泄漏
- 解决方案：使用 PyTorch 的图内存池，并在捕获后释放临时资源
动态控制流：
- 挑战：CUDA 图不支持动态控制流
- 解决方案：固定解码器路径，移除条件分支

CUDA 图捕获的优化策略

vLLM 采用了几个关键策略优化 CUDA 图性能：

分段捕获：

# 为模型不同部分捕获独立图
self.embedding_graph = self._capture_embedding_graph()
self.decoder_graphs = {}
for layer_idx in range(self.num_layers):
    self.decoder_graphs[layer_idx] = self._capture_decoder_layer(layer_idx)
self.final_layer_graph = self._capture_final_layer()

预热策略：

def warmup_model(self):
    """预热模型，确保图捕获前性能稳定"""
    # 使用典型批处理大小执行几次前向传播
    dummy_batch = self._create_dummy_batch(self.max_batch_size)
    for _ in range(3):  # 执行几轮预热
        self._forward_helper(**dummy_batch)
        torch.cuda.synchronize()

分层缓存：

def _get_graph_batch_sizes(self):
    """确定需要捕获的批处理大小集合"""
    # 使用分层策略，覆盖常见大小
    # 小批量：精确捕获
    # 中批量：对数间隔捕获
    # 大批量：线性间隔捕获
    sizes = []
    # 小批量（每个大小都捕获）
    sizes.extend(range(1, 17))
    # 中批量（对数间隔）
    sizes.extend([24, 32, 48, 64])
    # 大批量（线性间隔）
    sizes.extend([96, 128, 192, 256])
    return [s for s in sizes if s <= self.max_batch_size]

性能收益

CUDA 图捕获为 vLLM 带来显著性能提升：

吞吐量提升：
- 对于批量大小 8 的推理，吞吐量提高 15-30%
- 对于批量大小 32 的推理，吞吐量提高 25-40%
延迟降低：
- 单个请求延迟减少 8-15%
- 批处理请求平均延迟减少 20-35%
CPU 负载降低：
- CPU 使用率降低 30-50%，减轻调度瓶颈

使用限制与注意事项

尽管 CUDA 图提供了显著的性能优势，但也有一些限制需要注意：

固定输入形状：每个批处理大小需要单独的图
控制流限制：不支持动态分支和循环
内存消耗：维护多个图会增加 GPU 内存使用
调试困难：错误追踪更加复杂
硬件依赖：性能提升受 GPU 架构影响

实际应用示例

在 vLLM 的实际使用中，CUDA 图捕获的应用如下：

# 在 model_runner.py 的 execute_model 方法中
def execute_model(self, *args, **kwargs):
    """执行模型推理，自动选择最佳路径"""
    batch_size = kwargs["input_ids"].shape[0]
    
    # 选择执行路径
    if (self.model_captured and 
        batch_size in self.graph_runners and 
        not self.config.disable_cuda_graphs):
        # CUDA 图路径 - 高性能
        start_time = time.time()
        outputs = self._execute_with_graph(batch_size, *args, **kwargs)
        if self.profiler:
            self.profiler.record_cuda_graph_execution(time.time() - start_time)
        return outputs
    else:
        # 常规执行路径
        return self._execute_standard(*args, **kwargs)

实际的推理框架中，会自动选择使用 CUDA 图或标准执行路径，确保最佳性能。

总结而言，CUDA 图捕获是 vLLM 性能优化的关键组成部分，通过减少 CPU-GPU 交互开销，显著提高模型的推理吞吐量和降低延迟。这一技术特别适用于具有稳定形状和计算模式的大型语言模型推理场景。

5.4 KV 缓存管理

KV 缓存（Key-Value Cache）是大型语言模型自回归生成中的关键优化技术，它存储先前计算的注意力键值对，避免重复计算。vLLM 通过创新的缓存管理机制，使这一技术在大规模服务中更加高效。

KV 缓存基本原理

在 Transformer 的自回归解码中，每生成一个新 token，都需要计算该 token 与所有先前 token 的注意力：

# 未使用 KV 缓存的自回归生成
for i in range(max_seq_len):
    # 对所有已生成的 token 重新计算 K 和 V
    keys = self.compute_keys(input_ids[:, :i+1])    # 重复计算
    values = self.compute_values(input_ids[:, :i+1])  # 重复计算
    
    # 计算下一个 token 的输出
    next_token_logits = self.compute_attention(
        query=self.compute_query(input_ids[:, i:i+1]),
        keys=keys,
        values=values
    )
    
    # 采样下一个 token
    next_token = sample(next_token_logits)
    input_ids = torch.cat([input_ids, next_token], dim=1)

使用 KV 缓存后，先前计算的 K 和 V 被保留下来，只需为新 token 计算：

# 使用 KV 缓存的自回归生成
# 初始化缓存
key_cache = []
value_cache = []

for i in range(max_seq_len):
    if i == 0:
        # 处理提示词
        keys = self.compute_keys(input_ids)
        values = self.compute_values(input_ids)
        # 存入缓存
        key_cache.append(keys)
        value_cache.append(values)
    else:
        # 只计算最新 token 的 K 和 V
        new_key = self.compute_keys(input_ids[:, -1:])
        new_value = self.compute_values(input_ids[:, -1:])
        # 添加到缓存
        key_cache.append(new_key)
        value_cache.append(new_value)
    
    # 使用完整缓存计算注意力
    next_token_logits = self.compute_attention(
        query=self.compute_query(input_ids[:, -1:]),
        keys=torch.cat(key_cache, dim=1),
        values=torch.cat(value_cache, dim=1)
    )
    
    # 采样下一个 token
    next_token = sample(next_token_logits)
    input_ids = torch.cat([input_ids, next_token], dim=1)

这一优化将时间复杂度从 O(n²) 降低到 O(n)，其中 n 是序列长度。

vLLM 中的缓存结构

vLLM 使用层次化的结构管理 KV 缓存：

class KVCacheManager:
    def __init__(self, num_layers, num_heads, head_size, block_size):
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_size = head_size
        self.block_size = block_size
        
        # 物理缓存存储
        self.key_blocks = torch.empty(
            (0, block_size, num_heads, head_size),
            dtype=torch.float16, device="cuda"
        )
        self.value_blocks = torch.empty(
            (0, block_size, num_heads, head_size),
            dtype=torch.float16, device="cuda"
        )
        
        # 块分配器
        self.block_allocator = BlockAllocator()
        
        # 序列到块的映射
        self.block_tables = {}

vLLM 的 KV 缓存管理分为以下几个层级：

物理存储层：实际存储 K 和 V 的张量块
分配管理层：负责分配和回收物理块
逻辑映射层：维护序列 ID 到物理块的映射
访问控制层：提供高效的缓存读写接口

缓存分配与释放

缓存的分配和释放是 vLLM 内存管理的核心：

def allocate_kv_cache(self, seq_id: int, prompt_len: int) -> None:
    """为新序列分配 KV 缓存"""
    # 计算需要的块数
    num_blocks = ceil_div(prompt_len, self.block_size)
    
    # 分配物理块
    block_ids = []
    for _ in range(num_blocks):
        block_id = self.block_allocator.allocate()
        if block_id >= len(self.key_blocks):
            # 缓存扩容
            self._extend_cache(block_id + 1 - len(self.key_blocks))
        block_ids.append(block_id)
    
    # 更新块表
    self.block_tables[seq_id] = block_ids

def free_kv_cache(self, seq_id: int) -> None:
    """释放序列的 KV 缓存"""
    if seq_id not in self.block_tables:
        return
    
    # 获取分配的块
    block_ids = self.block_tables[seq_id]
    
    # 释放块
    for block_id in block_ids:
        self.block_allocator.free(block_id)
    
    # 移除映射
    del self.block_tables[seq_id]

缓存扩容是一个关键操作，它在需要更多物理块时动态调整缓存大小：

def _extend_cache(self, num_blocks: int) -> None:
    """扩展 KV 缓存"""
    current_size = len(self.key_blocks)
    new_size = current_size + num_blocks
    
    # 创建新存储空间
    new_key_blocks = torch.empty(
        (new_size, self.block_size, self.num_heads, self.head_size),
        dtype=torch.float16, device="cuda"
    )
    new_value_blocks = torch.empty(
        (new_size, self.block_size, self.num_heads, self.head_size),
        dtype=torch.float16, device="cuda"
    )
    
    # 复制现有数据
    if current_size > 0:
        new_key_blocks[:current_size] = self.key_blocks
        new_value_blocks[:current_size] = self.value_blocks
        
    # 替换缓存
    self.key_blocks = new_key_blocks
    self.value_blocks = new_value_blocks

缓存访问与更新

高效的缓存访问对性能至关重要：

def get_kv_cache(self, seq_id: int, positions: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """获取特定位置的 KV 缓存"""
    if seq_id not in self.block_tables:
        return None, None
    
    # 获取块表
    block_table = self.block_tables[seq_id]
    
    # 计算块索引和偏移
    block_indices = positions // self.block_size
    block_offsets = positions % self.block_size
    
    # 构建批索引张量
    indices_tensor = torch.empty(
        (len(positions), 3), dtype=torch.long, device="cuda"
    )
    
    # 填充索引 [block_idx, block_offset, head_idx]
    for i, (block_idx, offset) in enumerate(zip(block_indices, block_offsets)):
        indices_tensor[i, 0] = block_table[block_idx.item()]
        indices_tensor[i, 1] = offset
    
    # 提取缓存
    keys = torch.index_select(
        self.key_blocks.view(-1, self.num_heads, self.head_size),
        0, indices_tensor[:, 0] * self.block_size + indices_tensor[:, 1]
    )
    values = torch.index_select(
        self.value_blocks.view(-1, self.num_heads, self.head_size),
        0, indices_tensor[:, 0] * self.block_size + indices_tensor[:, 1]
    )
    
    return keys, values

在 vLLM 实际实现中，使用了更高效的 CUDA 自定义内核来加速缓存访问：

# CUDA 内核（伪代码）
@cuda.jit
def paged_attention_kernel(
    query,           # [batch_size, num_heads, head_size]
    key_cache,       # [num_blocks, block_size, num_heads, head_size]
    value_cache,     # [num_blocks, block_size, num_heads, head_size]
    block_tables,    # [batch_size, max_blocks_per_seq]
    seq_lens,        # [batch_size]
    output           # [batch_size, num_heads, head_size]
):
    # 获取线程索引
    batch_idx = cuda.blockIdx.x
    head_idx = cuda.blockIdx.y
    
    # 获取序列的块表和长度
    seq_block_table = block_tables[batch_idx]
    seq_len = seq_lens[batch_idx]
    
    # 获取当前查询向量
    q = query[batch_idx, head_idx]
    
    # 计算注意力分数和加权和
    acc = cuda.local.array(shape=(HEAD_SIZE,), dtype=float32)
    for i in range(HEAD_SIZE):
        acc[i] = 0.0
    
    # 遍历序列的所有块
    for block_idx in range(ceil_div(seq_len, BLOCK_SIZE)):
        physical_block = seq_block_table[block_idx]
        block_len = min(BLOCK_SIZE, seq_len - block_idx * BLOCK_SIZE)
        
        # 计算注意力分数
        scores = cuda.local.array(shape=(BLOCK_SIZE,), dtype=float32)
        for i in range(block_len):
            k = key_cache[physical_block, i, head_idx]
            scores[i] = dot_product(q, k)
        
        # 应用 softmax
        softmax_normalization(scores, block_len)
        
        # 计算加权和
        for i in range(block_len):
            v = value_cache[physical_block, i, head_idx]
            for j in range(HEAD_SIZE):
                acc[j] += scores[i] * v[j]
    
    # 写回输出
    for i in range(HEAD_SIZE):
        output[batch_idx, head_idx, i] = acc[i]

这个 CUDA 内核直接在 GPU 上执行分块注意力计算，避免了将分散的块复制到连续内存的开销。

序列复制与分叉

vLLM 支持序列复制和分叉操作，这对于并行采样和树搜索算法至关重要：

def fork_sequence(
    self, src_seq_id: int, dst_seq_id: int, fork_pos: int
) -> None:
    """在指定位置分叉序列，共享 KV 缓存"""
    if src_seq_id not in self.block_tables:
        return
    
    # 获取源序列的块表
    src_blocks = self.block_tables[src_seq_id]
    
    # 计算分叉位置对应的块
    fork_block_idx = fork_pos // self.block_size
    
    # 共享前 fork_block_idx 个块
    shared_blocks = src_blocks[:fork_block_idx + 1]
    
    # 复制源块引用
    if dst_seq_id in self.block_tables:
        # 清理可能的现有块
        self.free_kv_cache(dst_seq_id)
    
    # 创建新的块表
    self.block_tables[dst_seq_id] = shared_blocks.copy()
    
    # 增加共享块的引用计数
    for block_id in shared_blocks:
        self.block_allocator.increment_ref_count(block_id)

这种机制允许多个序列共享公共前缀的 KV 缓存，同时维护各自的后续缓存，大大节省了内存和计算资源。

缓存交换与重用

为了支持更多并发序列，vLLM 实现了缓存交换机制，将不活跃序列的缓存暂时移出 GPU：

def swap_out(self, seq_id: int) -> Dict:
    """将序列缓存交换出 GPU"""
    if seq_id not in self.block_tables:
        return None
    
    # 获取块表
    block_ids = self.block_tables[seq_id]
    
    # 准备 CPU 存储
    cpu_key_cache = []
    cpu_value_cache = []
    
    # 复制到 CPU
    for block_id in block_ids:
        cpu_key_cache.append(
            self.key_blocks[block_id].detach().cpu()
        )
        cpu_value_cache.append(
            self.value_blocks[block_id].detach().cpu()
        )
    
    # 释放 GPU 块
    self.free_kv_cache(seq_id)
    
    # 返回 CPU 缓存
    return {
        "block_table": block_ids,
        "key_cache": cpu_key_cache,
        "value_cache": cpu_value_cache
    }

def swap_in(self, seq_id: int, cpu_cache: Dict) -> None:
    """将序列缓存交换回 GPU"""
    # 提取 CPU 缓存
    block_table = cpu_cache["block_table"]
    cpu_key_cache = cpu_cache["key_cache"]
    cpu_value_cache = cpu_cache["value_cache"]
    
    # 分配新块
    new_block_ids = []
    for i in range(len(block_table)):
        # 分配块
        block_id = self.block_allocator.allocate()
        if block_id >= len(self.key_blocks):
            self._extend_cache(block_id + 1 - len(self.key_blocks))
        
        # 将数据复制回 GPU
        self.key_blocks[block_id].copy_(
            cpu_key_cache[i].to("cuda")
        )
        self.value_blocks[block_id].copy_(
            cpu_value_cache[i].to("cuda")
        )
        
        new_block_ids.append(block_id)
    
    # 更新块表
    self.block_tables[seq_id] = new_block_ids

这种交换机制使 vLLM 能够处理超出 GPU 内存容量的并发请求，通过智能调度，优先处理活跃序列。

优化策略

vLLM 在 KV 缓存管理中采用了多种优化策略：

预分配策略：

def preallocate_cache(self, max_seq_len: int, batch_size: int) -> None:
    """预分配 KV 缓存空间"""
    num_blocks_per_seq = ceil_div(max_seq_len, self.block_size)
    total_blocks = num_blocks_per_seq * batch_size
       
    if len(self.key_blocks) < total_blocks:
        self._extend_cache(total_blocks - len(self.key_blocks))

缓存压缩：降低精度以节省内存

def compress_cache(self) -> None:
    """将 KV 缓存压缩为更低精度"""
    if self.key_blocks.dtype != torch.float16:
        return
           
    # 将 float16 转换为 int8
    self.key_blocks_int8 = quantize(self.key_blocks)
    self.value_blocks_int8 = quantize(self.value_blocks)
       
    # 释放原始存储
    del self.key_blocks
    del self.value_blocks
       
    # 更新标志
    self.using_compressed_cache = True

层级缓存：不同层使用不同精度

def use_mixed_precision_cache(self) -> None:
    """对不同层使用不同精度缓存"""
    for i in range(self.num_layers):
        # 前几层使用低精度（对精度不太敏感）
        if i < self.num_layers // 3:
            self.layer_precisions[i] = "int8"
        # 中间层使用中等精度
        elif i < self.num_layers * 2 // 3:
            self.layer_precisions[i] = "fp16"
        # 后几层使用高精度（对精度更敏感）
        else:
            self.layer_precisions[i] = "fp32"

实际应用示例

在 vLLM 引擎的循环中，KV 缓存管理的应用如下：

def step(self):
    """执行一步生成"""
    # 1. 调度序列
    scheduled = self.scheduler.schedule()
    
    # 2. 准备批处理输入
    batch = self._prepare_batch(scheduled)
    
    # 3. 为新序列分配 KV 缓存
    for seq_id in batch.new_seq_ids:
        self.cache_manager.allocate_kv_cache(
            seq_id, batch.prompt_lens[seq_id]
        )
    
    # 4. 执行模型前向传播
    outputs = self.model_runner.forward(
        batch.input_ids,
        batch.positions,
        batch.kv_caches,
        batch.block_tables
    )
    
    # 5. 处理输出并更新序列状态
    for i, seq_id in enumerate(batch.seq_ids):
        # 采样下一个 token
        next_token = self._sample_next_token(outputs, i)
        
        # 添加到序列
        self.sequences[seq_id].append_token(next_token)
        
        # 检查是否完成
        if self._is_finished(seq_id):
            # 释放 KV 缓存
            self.cache_manager.free_kv_cache(seq_id)

总的来说，vLLM 的 KV 缓存管理系统结合了 PagedAttention 机制、动态内存分配和高效的 CUDA 内核，解决了大规模 LLM 服务中的内存效率和计算性能挑战。通过这些优化，vLLM 能够在有限的 GPU 内存中支持更多并发请求和更长的序列长度，为构建响应式的 LLM 服务提供了基础。

5.5 内存优化策略

vLLM 实现了多种内存优化策略，解决大型语言模型推理中的 GPU 内存瓶颈问题。这些策略从不同层面优化内存使用，提高服务的并发能力和稳定性。

内存压力来源

大型语言模型推理面临的主要内存压力来自三个方面：

模型权重：随着模型规模增长，权重参数占用大量内存

# 对于常见模型大小，权重内存需求
model_size_in_gb = num_parameters * dtype_size_in_bytes / (1024**3)
# 例如，70B 模型使用 FP16 精度需要约 140GB 内存

激活值：模型前向传播过程中的中间激活值

# 激活值内存需求（每个序列）
activation_size = num_layers * sequence_length * hidden_size * dtype_size

KV 缓存：自回归生成过程中存储的键值对

# KV 缓存内存需求（每个序列）
kv_cache_size = num_layers * sequence_length * 2 * hidden_size * dtype_size

vLLM 的内存优化策略针对这三个方面分别进行了优化。

张量并行与模型分片

为了处理超大模型，vLLM 实现了张量并行（Tensor Parallelism）技术：

class ModelSharding:
    def __init__(self, model, num_gpus):
        self.num_gpus = num_gpus
        self.model = model
        self.tp_size = num_gpus  # 张量并行度
        
    def shard_model(self):
        """将模型分片到多个 GPU"""
        # 1. 识别可并行层
        tp_modules = []
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Linear, nn.Embedding)):
                tp_modules.append((name, module))
        
        # 2. 为每个分片创建子模型
        self.devices = [f"cuda:{i}" for i in range(self.tp_size)]
        self.shards = [copy.deepcopy(self.model) for _ in range(self.tp_size)]
        
        # 3. 分割权重
        for name, module in tp_modules:
            if isinstance(module, nn.Linear):
                # 沿输出维度分割线性层
                self._shard_linear_layer(name)
            elif isinstance(module, nn.Embedding):
                # 沿词汇表维度分割嵌入层
                self._shard_embedding_layer(name)

张量并行的核心实现包括权重分割和梯度同步：

def _shard_linear_layer(self, name):
    """分割线性层权重"""
    original_weight = self.model.get_submodule(name).weight
    original_shape = original_weight.shape
    
    # 沿输出维度分割
    split_size = original_shape[0] // self.tp_size
    for i in range(self.tp_size):
        # 计算分片范围
        start_idx = i * split_size
        end_idx = (i + 1) * split_size if i < self.tp_size - 1 else original_shape[0]
        
        # 分配分片
        shard = self.shards[i].get_submodule(name)
        with torch.no_grad():
            shard.weight = nn.Parameter(
                original_weight[start_idx:end_idx].clone().to(self.devices[i])
            )
        
        # 更新形状
        shard.in_features = original_shape[1]
        shard.out_features = end_idx - start_idx

vLLM 使用自定义的并行前向传播来协调多 GPU 计算：

def tp_forward(self, input_tensor):
    """使用张量并行执行前向传播"""
    # 1. 分发输入到各个设备
    input_shards = []
    for i in range(self.tp_size):
        input_shards.append(input_tensor.to(self.devices[i]))
    
    # 2. 并行执行每个分片
    output_shards = []
    for i in range(self.tp_size):
        with torch.cuda.device(self.devices[i]):
            output_shards.append(self.shards[i](input_shards[i]))
    
    # 3. 合并结果
    if isinstance(output_shards[0], torch.Tensor):
        # 对于分割的线性层，沿输出维度合并
        return torch.cat(output_shards, dim=-1)
    else:
        # 对于其他层，可能需要特殊处理
        return self._merge_custom_outputs(output_shards)

量化与混合精度

vLLM 支持多种量化技术以降低内存需求：

def load_quantized_model(model_path, quantization="int8"):
    """加载量化模型"""
    if quantization == "int8":
        # 使用 INT8 量化
        return load_int8_model(model_path)
    elif quantization == "int4":
        # 使用 INT4 量化
        return load_int4_model(model_path)
    elif quantization == "awq":
        # 使用 AWQ 量化
        return load_awq_model(model_path)
    else:
        # 默认加载全精度模型
        return load_full_precision_model(model_path)

量化策略具体实现示例：

def load_int8_model(model_path):
    """加载 INT8 量化模型"""
    # 1. 加载原始模型配置和权重
    config = AutoConfig.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_config(config)
    state_dict = torch.load(f"{model_path}/pytorch_model.bin")
    
    # 2. 对每个权重进行量化
    quantized_state_dict = {}
    for name, param in state_dict.items():
        # 跳过某些不应量化的层
        if "ln" in name or "embedding" in name:
            quantized_state_dict[name] = param
        else:
            # 执行 INT8 量化
            quantized_param = quantize_to_int8(param)
            quantized_state_dict[name] = quantized_param
    
    # 3. 加载量化后的权重
    model.load_state_dict(quantized_state_dict)
    
    # 4. 注册自定义内核用于量化计算
    replace_with_quantized_kernels(model)
    
    return model

vLLM 还实现了更复杂的二阶量化和混合精度策略：

class MixedPrecisionManager:
    def __init__(self, model):
        self.model = model
        
    def apply_mixed_precision(self):
        """应用混合精度策略"""
        # 分析层的重要性
        layer_importance = self._analyze_layer_importance()
        
        # 应用不同精度
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Linear):
                layer_name = name.split(".")[-2]  # 获取层名称
                if layer_name in layer_importance:
                    importance = layer_importance[layer_name]
                    
                    # 根据重要性选择精度
                    if importance > 0.8:
                        # 重要层使用更高精度
                        module.to(torch.float16)
                    elif importance > 0.5:
                        # 中等重要性使用 INT8
                        quantize_module_to_int8(module)
                    else:
                        # 低重要性使用 INT4
                        quantize_module_to_int4(module)

激活检查点

对于内存受限的场景，vLLM 实现了激活检查点（Activation Checkpointing）技术：

class ActivationCheckpointing:
    def __init__(self, model):
        self.model = model
        
    def apply_checkpointing(self):
        """对模型应用激活检查点"""
        for name, module in self.model.named_children():
            if isinstance(module, TransformerLayer):
                # 包装 Transformer 层使用检查点
                self.model._modules[name] = torch.utils.checkpoint.checkpoint_wrapper(module)

实现自定义粒度的检查点，以平衡内存使用和计算开销：

def custom_checkpointing(self):
    """实现自定义粒度的检查点"""
    checkpoint_layers = []
    
    # 为每个 Transformer 层选择检查点策略
    for i, layer in enumerate(self.model.layers):
        if i % 2 == 0:  # 偶数层：细粒度检查点
            # 单独对 attention 和 mlp 应用检查点
            layer.attention = torch.utils.checkpoint.checkpoint_wrapper(
                layer.attention, 
                preserve_rng_state=False
            )
            layer.mlp = torch.utils.checkpoint.checkpoint_wrapper(
                layer.mlp,
                preserve_rng_state=False
            )
        else:  # 奇数层：整层检查点
            checkpoint_layers.append(layer)
    
    # 对选择的层应用检查点
    for layer in checkpoint_layers:
        layer_index = self.model.layers.index(layer)
        self.model.layers[layer_index] = torch.utils.checkpoint.checkpoint_wrapper(
            layer, 
            preserve_rng_state=False
        )

优化的内存管理器

vLLM 实现了自定义的内存管理器，更精细地控制内存分配和释放：

class GPUMemoryManager:
    def __init__(self, device, max_memory_fraction=0.9):
        self.device = device
        self.max_memory = int(torch.cuda.get_device_properties(device).total_memory 
                              * max_memory_fraction)
        self.allocated = 0
        self.buffers = {}
        self.free_buffers = {}
        
    def allocate(self, size, dtype):
        """分配内存缓冲区"""
        # 检查现有的空闲缓冲区
        buffer_key = (size, dtype)
        if buffer_key in self.free_buffers and self.free_buffers[buffer_key]:
            # 重用现有缓冲区
            buffer = self.free_buffers[buffer_key].pop()
            return buffer
        
        # 检查内存限制
        bytes_size = size * dtype_size(dtype)
        if self.allocated + bytes_size > self.max_memory:
            # 尝试回收内存
            self._garbage_collect()
            if self.allocated + bytes_size > self.max_memory:
                raise RuntimeError(f"Out of memory: trying to allocate {bytes_size} bytes")
            
            # 分配新缓冲区
            buffer = torch.empty(size, dtype=dtype, device=self.device)
            self.allocated += bytes_size
            buffer_id = id(buffer)
            self.buffers[buffer_id] = (buffer, size, dtype, bytes_size)
            
            return buffer
            
        def free(self, buffer):
            """释放缓冲区（移至空闲池）"""
            buffer_id = id(buffer)
            if buffer_id in self.buffers:
                _, size, dtype, _ = self.buffers[buffer_id]
                buffer_key = (size, dtype)
                
                # 添加到空闲池
                if buffer_key not in self.free_buffers:
                    self.free_buffers[buffer_key] = []
                self.free_buffers[buffer_key].append(buffer)
        
        def _garbage_collect(self):
            """强制垃圾回收释放内存"""
            # 清除超过阈值的空闲缓冲区
            for buffer_key in list(self.free_buffers.keys()):
                if len(self.free_buffers[buffer_key]) > 5:  # 保留一定数量的缓冲区
                    # 释放多余缓冲区
                    excess_buffers = self.free_buffers[buffer_key][5:]
                    self.free_buffers[buffer_key] = self.free_buffers[buffer_key][:5]
                    
                    # 更新分配跟踪
                    for buffer in excess_buffers:
                        buffer_id = id(buffer)
                        if buffer_id in self.buffers:
                            _, _, _, bytes_size = self.buffers[buffer_id]
                            self.allocated -= bytes_size
                            del self.buffers[buffer_id]
                    
                    # 明确删除引用
                    del excess_buffers
            
            # 强制 PyTorch 缓存清理
            torch.cuda.empty_cache()

智能的内存使用监控

vLLM 实现了内存使用监控，动态调整行为：

class MemoryMonitor:
    def __init__(self, device, warning_threshold=0.9, critical_threshold=0.95):
        self.device = device
        self.warning_threshold = warning_threshold
        self.critical_threshold = critical_threshold
        self.last_check_time = 0
        self.check_interval = 1.0  # 秒
        
    def check_memory(self):
        """检查内存使用情况并作出响应"""
        current_time = time.time()
        if current_time - self.last_check_time < self.check_interval:
            return MemoryStatus.OK
            
        # 更新检查时间
        self.last_check_time = current_time
        
        # 获取当前内存使用
        used_memory = torch.cuda.memory_allocated(self.device)
        total_memory = torch.cuda.get_device_properties(self.device).total_memory
        memory_fraction = used_memory / total_memory
        
        # 根据使用率确定状态
        if memory_fraction > self.critical_threshold:
            return MemoryStatus.CRITICAL
        elif memory_fraction > self.warning_threshold:
            return MemoryStatus.WARNING
        else:
            return MemoryStatus.OK
    
    def get_recommended_actions(self, status):
        """根据内存状态提供建议操作"""
        if status == MemoryStatus.CRITICAL:
            return [
                "swap_out_inactive_sequences",
                "reduce_batch_size",
                "force_garbage_collect"
            ]
        elif status == MemoryStatus.WARNING:
            return [
                "delay_new_requests",
                "compress_kv_cache"
            ]
        else:
            return []

这个监控器与调度器集成，动态调整系统行为：

def adaptive_sequence_scheduling(self):
    """根据内存使用自适应调整调度"""
    # 检查内存状态
    memory_status = self.memory_monitor.check_memory()
    actions = self.memory_monitor.get_recommended_actions(memory_status)
    
    # 应用推荐操作
    if "reduce_batch_size" in actions:
        # 临时降低批大小
        self.current_max_batch_size = max(1, self.max_batch_size // 2)
    else:
        # 恢复正常批大小
        self.current_max_batch_size = self.max_batch_size
        
    if "swap_out_inactive_sequences" in actions:
        # 交换出不活跃序列
        inactive_seqs = self._identify_inactive_sequences()
        for seq_id in inactive_seqs:
            self.swap_out_sequence(seq_id)
            
    if "force_garbage_collect" in actions:
        # 强制垃圾回收
        torch.cuda.empty_cache()
        gc.collect()

实际应用案例与性能收益

通过这些内存优化策略的组合应用，vLLM 在实际部署中取得了显著的性能改进：

并发能力提升：

# 应用优化前（传统方法）
max_concurrent_requests = 5  # 13B 模型，A100 GPU
   
# 应用优化后（vLLM）
max_concurrent_requests = 20  # 同样配置下增加4倍

内存效率提升：

# 传统部署（静态批处理）
memory_per_request = 8 * hidden_size * max_seq_len * num_layers
   
# vLLM（分页注意力+内存优化）
memory_per_request = 8 * hidden_size * actual_seq_len * num_layers / efficiency_factor
# efficiency_factor ≈ 2.5-3.0

量化效益：

# FP16 模型
model_size_gb = 13  # 对于13B模型
   
# INT8 量化后
model_size_gb = 6.5  # 节省50%内存
   
# INT4 量化后
model_size_gb = 3.25  # 节省75%内存

总的来说，vLLM 的内存优化策略在多个层面协同工作，从模型权重优化、中间激活管理到 KV 缓存组织，全方位提升了大型语言模型的推理效率。这使得在单个 GPU 上可以部署更大的模型或服务更多的并发用户，大幅降低了 LLM 服务的硬件需求和运营成本。

5.6 并行计算优化

vLLM 实现了多种并行计算优化策略，提高服务性能和并发能力。这些策略从不同层面优化计算效率，减少资源消耗，提高服务质量。

并行计算的基本原理

并行计算的核心思想是将一个任务分解为多个子任务，并在多个处理器上同时执行这些子任务。vLLM 通过以下几种方式实现并行计算：

多线程：利用操作系统提供的线程机制，在单个进程中同时执行多个线程。
多进程：利用操作系统提供的进程机制，在单个机器上同时运行多个进程。
分布式计算：利用网络通信机制，在多台机器上同时执行计算任务。

vLLM 的并行计算策略包括：

线程池：预先创建一组线程，用于处理并发请求。
进程池：预先创建一组进程，用于处理并发请求。
分布式调度：利用分布式计算框架（如 Ray）实现任务调度。

线程池与进程池

vLLM 实现了线程池和进程池机制，用于管理并发线程和进程。这些机制可以显著提高服务性能，减少线程切换和进程切换的开销。

# 线程池示例
thread_pool = ThreadPool(max_workers=10)

# 进程池示例
process_pool = Pool(processes=4)

分布式调度

vLLM 利用分布式计算框架（如 Ray）实现分布式调度，将任务分配到多个处理器上执行。分布式调度可以显著提高服务性能，减少单点故障的风险。

# Ray 示例
from ray import serve

@serve.deployment
class MyModel:
    def __call__(self, request):
        # 分布式任务调度
        result = ray.get(ray.remote(compute_task)(request))
        return result

实际应用案例

在 vLLM 中，并行计算优化策略的应用如下：

线程池：在每次前向传播中创建多个线程，用于处理并发请求。

# 线程池示例
thread_pool = ThreadPool(max_workers=10)
   
# 在每次前向传播中创建线程
def execute_model(self, input_ids, positions, kv_caches, block_tables):
    # 创建线程池
    with ThreadPoolExecutor(max_workers=10) as executor:
        futures = [executor.submit(self._forward_helper, input_ids, positions, kv_caches, block_tables) for _ in range(10)]
        results = [future.result() for future in futures]
    return results

进程池：在每次前向传播中创建多个进程，用于处理并发请求。

# 进程池示例
process_pool = Pool(processes=4)
   
# 在每次前向传播中创建进程
def execute_model(self, input_ids, positions, kv_caches, block_tables):
    # 创建进程池
    with Pool(processes=4) as pool:
        results = pool.map(self._forward_helper, [input_ids] * 4, [positions] * 4, [kv_caches] * 4, [block_tables] * 4)
    return results

分布式调度：利用分布式计算框架（如 Ray）实现分布式调度，将任务分配到多个处理器上执行。

# Ray 示例
from ray import serve

@serve.deployment
class MyModel:
    def __call__(self, request):
        # 分布式任务调度
        result = ray.get(ray.remote(compute_task)(request))
        return result

总的来说，vLLM 的并行计算优化策略显著提高了服务性能和并发能力，减少了资源消耗，提高了服务质量。这些策略从不同层面优化计算效率，减少线程切换和进程切换的开销，同时提高了系统的稳定性和可靠性。

5. 架构与类关系图

5.1 整体架构概览

vLLM的架构设计采用了模块化和层次化的方式，清晰分离了不同功能组件，使系统具有高度的可扩展性和可维护性。下图展示了vLLM的整体架构:

+---------------------+
|     用户 API 层      |  (LLM, AsyncLLM 接口)
+----------+----------+
           |
+----------v----------+
|    LLMEngine 层     |  (核心协调器)
+----------+----------+
           |
     +-----+-----+
     |           |
+----v----+ +----v----+
|Scheduler| |Worker池 |  (调度与执行)
+----+----+ +----+----+
     |           |
     |      +----v----+
     |      |Model    |  (模型执行)
     |      |Runner   |
     |      +----+----+
     |           |
+----v----+ +----v----+
|请求队列 | |BlockMgr |  (资源管理)
+----+----+ +----+----+
     |           |
+----v-----------v----+
|     底层优化层       |  (CUDA图, KV缓存)
+---------------------+

5.2 核心类结构图

vLLM的主要类及其关系如下：

LLMEngine
  ├── ModelConfig         // 模型配置
  ├── ParallelConfig      // 并行策略配置
  ├── SchedulerConfig     // 调度器配置
  ├── CacheConfig         // 缓存配置
  │
  ├── Tokenizer           // 分词器
  ├── RequestTracker      // 请求跟踪器
  │    └── Request        // 请求对象
  │
  ├── Scheduler[]         // 调度器数组
  │    ├── SequenceGroup  // 序列组
  │    └── Sequence       // 序列
  │
  └── WorkerPool          // Worker池
       ├── Worker         // Worker实例
       │    ├── ModelRunner        // 模型执行器
       │    │    ├── AttentionState      // 注意力状态
       │    │    └── CUDAGraphRunner     // CUDA图执行器
       │    │
       │    └── BlockSpaceManager  // 块空间管理器
       │         ├── BlockAllocator      // 块分配器
       │         └── BlockTable          // 块表
       │
       └── ModelReplica   // 模型副本

5.3 关键组件交互流程

vLLM的核心交互流程如下图所示，展示了请求从接收到处理完成的整个生命周期：

Client → [LLM API] → LLMEngine.add_request()
  │                         │
  │                         ▼
  │                  RequestTracker
  │                         │
  │                         ▼
  │                  Scheduler.add_request()
  │                         │
  │                         ▼
  │                  Engine Core Loop
  │                         │
  │                         ▼
  │                  Scheduler.schedule()
  │                         │
  │                         ▼
  │                  Worker.execute_model()
  │                         │
  │                         ▼
  │                  ModelRunner.forward()
  │                         │
  │                         ▼
Client ← [Results] ← Engine._process_outputs()

5.4 并行与分布式架构

vLLM支持多种并行策略，包括张量并行、流水线并行和数据并行：

数据并行 (多副本)
+----------+  +----------+
| Engine 1 |  | Engine 2 |
+-----+----+  +----+-----+
      |            |
+-----v----+  +----v-----+
| Model 1  |  | Model 2  |
+----------+  +----------+

张量并行 (单模型分片)
+-------------------+
|      Engine       |
+---------+---------+
          |
+---------v---------+
| 分布式模型 (TP=2) |
+---+-------------+-+
    |             |
+---v---+     +---v---+
|GPU 0  |     |GPU 1  |
|模型分片|     |模型分片|
+-------+     +-------+

5.5 内存管理架构

vLLM的PagedAttention内存管理模型：

序列1    序列2    序列3    ← 逻辑序列
 ↓        ↓        ↓
+---+    +---+    +---+
| A |    | C |    | E |    ← 块表
+---+    +---+    +---+
 ↓        ↓        ↓
+---+    +---+    +---+
| 1 |    | 3 |    | 5 |    ← 物理块索引
+---+    +---+    +---+
 ↓        ↓        ↓
+---+    +---+    +---+
| B |    | D |    | F |    ← 块表
+---+    +---+    +---+
 ↓        ↓        ↓
+---+    +---+    +---+
| 2 |    | 4 |    | 6 |    ← 物理块索引
+---+    +---+    +---+

物理内存:
+---+---+---+---+---+---+
| 1 | 2 | 3 | 4 | 5 | 6 |    ← 物理块
+---+---+---+---+---+---+

这种分页架构允许高效利用GPU内存，使不同长度的序列能够共享物理内存，大幅提升内存利用率。

6. 高级 Python 语法应用

vLLM 代码中应用了多种高级 Python 语法特性，这些特性提高了代码的灵活性、可维护性和性能。即使你不熟悉vLLM，理解这些Python高级特性也能帮助你写出更好的代码。

6.1 类型注解与泛型编程

vLLM 广泛使用类型注解提高代码可读性和健壮性：

# 类型变量的使用
ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner

这种语法表示 ModelRunnerClass 是一个类型变量，它指向的类必须是 GPUModelRunnerBase 的子类。这不仅提供了类型安全性，还使代码更加自解释。

类型注解的优势：

提供更好的 IDE 支持和自动补全
在编译时捕获类型错误
使接口约束更明确
支持泛型编程模式

vLLM 中常见的类型注解模式：

from typing import Dict, List, Optional, Tuple, Union, Type, TypeVar, Generic

# 基本类型注解
def process_tokens(tokens: List[int]) -> List[str]:
    """处理token IDs并返回字符串列表"""
    return [str(token) for token in tokens]

# 泛型和类型变量
T = TypeVar('T', bound=BaseModel)
class Registry(Generic[T]):
    """通用模型注册表"""
    def __init__(self):
        self.models: Dict[str, T] = {}
        
    def register(self, model: T) -> None:
        self.models[model.name] = model
        
    def get(self, name: str) -> Optional[T]:
        return self.models.get(name)

# 可选类型
def get_cache(key: str) -> Optional[torch.Tensor]:
    """从缓存获取张量，可能返回None"""
    if key in cache:
        return cache[key]
    return None

# 联合类型
def process_input(input_data: Union[str, List[str], Dict[str, str]]) -> None:
    """处理多种可能的输入类型"""
    if isinstance(input_data, str):
        # 处理单个字符串
        pass
    elif isinstance(input_data, list):
        # 处理字符串列表
        pass
    else:
        # 处理字典
        pass

在vLLM核心组件如引擎和调度器中，类型注解扮演着关键角色。例如，引擎的参数定义：

class EngineArgs:
    """引擎参数配置"""
    
    def __init__(
        self,
        model: str,
        tokenizer: Optional[str] = None,
        tensor_parallel_size: int = 1,
        dtype: Optional[str] = None,
        quantization: Optional[str] = None,
        revision: Optional[str] = None,
        tokenizer_revision: Optional[str] = None,
        seed: int = 0,
        max_model_len: Optional[int] = None,
        worker_use_ray: bool = False,
        pipeline_parallel_size: int = 1,
        block_size: int = 16,
        swap_space: int = 4,
        gpu_memory_utilization: float = 0.9,
        max_num_batched_tokens: Optional[int] = None,
        max_num_seqs: int = 256,
        max_paddings: int = 256,
        disable_log_stats: bool = False,
        revision_date: Optional[str] = None,
        tokenizer_mode: str = "auto",
        trust_remote_code: bool = False,
        max_log_len: Optional[int] = None,
        tokenizer_timeout_s: int = 60,
    ) -> None:
        ...

这种详细的类型注解使得参数含义一目了然，同时也在运行时提供类型检查。通过 mypy 等静态类型检查工具，可以在代码运行前发现潜在的类型错误。

在工作原理上，Python 的类型注解是通过 PEP 484、PEP 526 等规范实现的，允许开发者在代码中标注变量、参数和返回值的类型。这些注解在运行时不会强制执行类型检查，但它们可以被 IDE 和静态分析工具使用。

vLLM 项目通过严格的类型注解，显著提高了代码质量和开发效率，特别是在处理复杂的并行计算和内存管理时，类型安全对于避免潜在错误至关重要。

6.2 动态类解析与反射机制

vLLM 使用 Python 的反射机制实现动态类加载和实例化，这对于构建灵活的架构至关重要：

def resolve_class_by_name(class_name: str) -> Type:
    """根据完全限定名解析类"""
    module_name, class_name = class_name.rsplit(".", 1)
    module = importlib.import_module(module_name)
    return getattr(module, class_name)

# 使用示例
worker_class_name = config.get("worker_class", "vllm.worker.Worker")
worker_class = resolve_class_by_name(worker_class_name)
worker = worker_class(**worker_args)

这种动态加载机制的关键实现细节：

模块导入：使用 importlib.import_module 在运行时导入模块
属性获取：使用 getattr 动态获取模块中的类或函数
运行时实例化：根据字符串配置动态创建类实例

在 vLLM 中，这种机制广泛应用于多个组件：

# 在 worker.py 中动态加载模型实现
def _initialize_worker(model_config):
    model_type = model_config.model_type
    
    if model_type == "gpt_neox":
        model_cls_name = "vllm.model.GPTNeoXForCausalLM"
    elif model_type == "llama":
        model_cls_name = "vllm.model.LlamaForCausalLM"
    elif model_type == "mistral":
        model_cls_name = "vllm.model.MistralForCausalLM"
    else:
        # 支持完全自定义模型
        model_cls_name = model_config.get("model_class")
    
    # 动态加载模型类
    model_cls = resolve_class_by_name(model_cls_name)
    return model_cls(model_config)

vLLM 的执行器系统也依赖于动态类解析：

def get_executor(args: EngineArgs) -> ExecutorBase:
    """根据配置创建合适的执行器"""
    if args.distributed_executor_backend:
        # 根据配置选择不同的分布式执行器
        if args.distributed_executor_backend == "ray":
            executor_cls_name = "vllm.executor.RayExecutor"
        elif args.distributed_executor_backend == "mp":
            executor_cls_name = "vllm.executor.MultiprocessingExecutor"
        else:
            executor_cls_name = args.distributed_executor_backend
            
        # 动态加载执行器类
        executor_cls = resolve_class_by_name(executor_cls_name)
        return executor_cls(args)
    else:
        # 默认使用本地执行器
        return LocalExecutor(args)

除了类加载，vLLM 还使用反射来实现参数检查和默认值处理：

def create_engine_from_args(args_dict: Dict[str, Any]) -> LLMEngine:
    """从字典创建引擎实例"""
    # 获取 EngineArgs 的签名
    sig = inspect.signature(EngineArgs.__init__)
    
    # 过滤有效参数
    valid_params = {}
    for param_name, param in sig.parameters.items():
        if param_name in args_dict:
            valid_params[param_name] = args_dict[param_name]
            
    # 创建 EngineArgs 实例
    engine_args = EngineArgs(**valid_params)
    
    # 创建引擎
    return LLMEngine.from_engine_args(engine_args)

这种动态类解析与反射机制的优势：

插件架构支持：允许不修改核心代码的情况下扩展功能
配置驱动的组件加载：通过配置文件或环境变量控制实现细节
避免硬编码依赖：降低组件间耦合，提高代码可维护性
运行时组件替换：支持根据环境和需求动态选择最佳实现

作为一个具体例子，用户可以通过配置指定一个自定义的调度器类：

# 用户代码
custom_args = {
    "model": "llama-3.2-1b-instruct",
    "scheduler_class": "my_project.schedulers.CustomPriorityScheduler"
}
engine = create_engine_from_args(custom_args)

内部实现会动态加载用户提供的调度器类并使用它，而不需要修改 vLLM 的核心代码。这种灵活性使 vLLM 能够适应各种特定需求，同时保持核心代码的稳定性。

总的来说，动态类解析和反射机制是 vLLM 灵活架构的关键，它使得系统能够根据需求动态加载不同组件，适应不同的硬件环境和使用场景。

6.3 方法委托与魔术方法

vLLM 使用 __getattr__ 魔术方法实现方法委托，这是一种强大的设计模式：

class WorkerWrapperBase:
    def __init__(self, worker):
        self.worker = worker
        
    def __getattr__(self, attr):
        return getattr(self.worker, attr)

这段看似简单的代码实际上实现了一个强大的代理模式。当在 WorkerWrapperBase 实例上访问一个不存在的属性或方法时，Python 会调用 __getattr__ 方法。这个方法将请求委托给内部的 worker 对象，实现了透明的方法转发。

具体工作流程：

客户代码调用 worker_wrapper.some_method()
如果 worker_wrapper 没有 some_method 属性，Python 调用 __getattr__('some_method')
__getattr__ 返回 worker.some_method
客户代码实际调用了 worker.some_method()

在 vLLM 中，这种模式的应用场景：

# 在 executor.py 中
class LocalExecutor(ExecutorBase):
    def __init__(self, args):
        # 创建本地 Worker
        self.worker = Worker(...)
    
    def __getattr__(self, attr):
        # 将方法调用委托给内部 Worker
        return getattr(self.worker, attr)
    
# 在 engine.py 中
def init_device(self):
    # LocalExecutor 没有 init_device 方法，但这个调用会被委托给内部的 worker
    self.executor.init_device()

另一个例子是 vLLM 的 DelegateWorkerBase 类：

class DelegateWorkerBase(WorkerBase):
    def __init__(self, worker: WorkerBase):
        self.worker = worker
    
    def __getattr__(self, attr):
        return getattr(self.worker, attr)

这种设计模式的优势：

代码减少：不需要为每个要转发的方法编写样板代码
透明代理：客户代码无需了解委托机制
动态性：支持运行时添加的方法，无需修改代理类
装饰器模式：可以轻松在方法调用前后添加额外行为

例如，添加日志记录功能：

class LoggingWorkerWrapper(WorkerWrapperBase):
    def __getattr__(self, attr):
        original_attr = super().__getattr__(attr)
        if callable(original_attr):
            def wrapper(*args, **kwargs):
                logger.debug(f"Calling {attr} with args: {args}, kwargs: {kwargs}")
                result = original_attr(*args, **kwargs)
                logger.debug(f"{attr} returned: {result}")
                return result
            return wrapper
        return original_attr

在 vLLM 中，这种方法委托模式解决了一个关键问题：WorkerBase 定义了许多抽象方法，如 init_device 和 load_model，这些方法在子类中必须实现。通过方法委托，可以创建特殊的 Worker 包装器，而不必实现所有抽象方法，只需将它们委托给内部的 Worker 实例。

# 抽象基类定义
class WorkerBase:
    def init_device(self):
        raise NotImplementedError
        
    def load_model(self):
        raise NotImplementedError

# 具体实现
class CUDAWorker(WorkerBase):
    def init_device(self):
        # CUDA 实现
        ...
    
    def load_model(self):
        # CUDA 模型加载
        ...

# 代理类 - 无需实现抽象方法
class MonitoringWorker(WorkerBase):
    def __init__(self, worker: WorkerBase):
        self.worker = worker
        self.metrics = {}
    
    def __getattr__(self, attr):
        return getattr(self.worker, attr)
    
    # 只添加新功能，无需重新实现所有方法
    def collect_metrics(self):
        # 收集性能指标
        ...

通过魔术方法实现的方法委托是 vLLM 灵活架构的关键组成部分，它减少了代码重复，提高了组件间的解耦性，并使系统更容易扩展。

6.4 弱引用与特殊调用模式

vLLM 在处理对象引用时使用了 weakref.proxy，特别是在注意力状态初始化中：

class ModelRunner:
    def initialize_attention_state(self):
        # 创建注意力状态类的实例
        self.attn_state = self.attn_backend.get_state_cls()(weakref.proxy(self))

这行代码结合了三种高级 Python 特性：

动态类获取：get_state_cls() 返回类定义而非实例
弱引用：weakref.proxy(self) 创建对象的弱引用
链式调用：立即实例化获取的类

让我们详细分析这些特性：

动态类获取

get_state_cls() 方法返回一个类定义，而不是类的实例：

def get_state_cls(self):
    """返回适用于当前后端的注意力状态类"""
    if self.use_flash_attention:
        return FlashAttentionState
    elif self.use_varlen_attention:
        return VarLenAttentionState
    else:
        return StandardAttentionState

这种工厂方法模式允许系统根据运行时条件选择适当的实现。

弱引用与 weakref.proxy

weakref.proxy(self) 创建了一个代理对象，它的行为类似于对原始对象的引用，但不增加引用计数：

# 常规引用
runner = ModelRunner()
state = AttentionState(runner)  # runner 的引用计数加 1

# 弱引用
runner = ModelRunner()
state = AttentionState(weakref.proxy(runner))  # runner 的引用计数不变

弱引用的主要优势：

避免循环引用：ModelRunner 包含 attn_state，而 attn_state 需要引用回 ModelRunner。如果使用普通引用，会形成循环引用，导致内存泄漏。
不阻止垃圾回收：当所有普通引用消失时，对象可以被回收，即使仍有弱引用。
自动失效：当被引用对象被回收时，weakref.proxy 会抛出 ReferenceError，而不是悄悄返回陈旧数据。

在大型模型处理中，内存管理是核心挑战，弱引用帮助 vLLM 管理复杂的对象关系，避免不必要的内存消耗。

链式调用模式

链式调用是将多个操作组合成一行代码的技术：

# 拆分写法
cls = self.attn_backend.get_state_cls()
self.attn_state = cls(weakref.proxy(self))

# 链式写法
self.attn_state = self.attn_backend.get_state_cls()(weakref.proxy(self))

这种模式提高了代码简洁性，减少了中间变量，同时保持了可读性。

实际应用案例

在 vLLM 中，注意力状态类需要访问 ModelRunner 的方法和属性，但同时 ModelRunner 也需要持有注意力状态实例。这是一个典型的循环引用场景：

class AttentionState:
    def __init__(self, model_runner):
        self.model_runner = model_runner
        
    def compute_attention(self, query):
        # 访问 model_runner 的属性和方法
        batch_size = self.model_runner.batch_size
        ...

class ModelRunner:
    def initialize(self):
        # 创建并持有注意力状态
        self.attn_state = AttentionState(self)  # 直接引用会造成循环引用

使用弱引用解决这个问题：

class AttentionState:
    def __init__(self, model_runner):
        # 使用弱引用避免循环引用
        self.model_runner = weakref.proxy(model_runner)

class ModelRunner:
    def initialize(self):
        # 安全地创建注意力状态
        self.attn_state = AttentionState(self)

这种模式在大型模型中特别重要，因为这些模型处理大量内存，任何内存泄漏都可能导致严重问题。

此外，vLLM 还在缓存管理和事件回调系统中使用了弱引用：

class CacheManager:
    def __init__(self):
        # 使用 WeakValueDictionary 避免阻止对象回收
        self.cache = weakref.WeakValueDictionary()
    
    def add_to_cache(self, key, value):
        self.cache[key] = value
    
    def get_from_cache(self, key):
        return self.cache.get(key)

总结起来，vLLM 通过组合动态类获取、弱引用和链式调用，实现了高效、灵活且内存安全的对象交互模式，这对于管理大型语言模型的复杂组件关系至关重要。

6.5 类方法与工厂模式

vLLM 使用类方法实现工厂模式，这是一种创建对象的优雅方式：

class LLMEngine:
    @classmethod
    def from_engine_args(cls, args: EngineArgs) -> "LLMEngine":
        # 创建并配置引擎实例
        engine = cls()
        executor = get_executor(args)
        executor.init_workers(args)
        executor.init_device()
        executor.load_model()
        return engine

类方法与普通方法和静态方法的区别在于它接收类本身作为第一个参数（通常命名为 cls），而不是实例（self）。这使得类方法可以访问和操作类级别的属性和方法，同时也能创建类的实例。

工厂模式实现

在 vLLM 中，类方法作为工厂函数非常常见：

class ModelLoader:
    @classmethod
    def from_pretrained(cls, model_path: str, **kwargs) -> "ModelLoader":
        # 1. 检查模型类型
        model_type = cls._detect_model_type(model_path)
        
        # 2. 加载配置
        config = cls._load_config(model_path, model_type)
        
        # 3. 准备参数
        params = {**config, **kwargs}
        
        # 4. 创建适当的实例
        return cls(params)

这种方法的优势包括：

封装实例化逻辑：隐藏复杂的创建过程
提供友好的接口：比直接调用构造函数更易用
支持多态：子类可以重写和扩展工厂方法
预处理参数：在创建实例前验证和准备参数

多态行为

类方法与继承配合使用时特别强大：

class BaseScheduler:
    @classmethod
    def create(cls, config: Dict[str, Any]) -> "BaseScheduler":
        return cls(config)

class FIFOScheduler(BaseScheduler):
    pass
    
class PriorityScheduler(BaseScheduler):
    @classmethod
    def create(cls, config: Dict[str, Any]) -> "PriorityScheduler":
        # 特殊处理优先级配置
        if "priorities" not in config:
            config["priorities"] = {"default": 1}
        return super().create(config)

# 使用示例
schedulers = {
    "fifo": FIFOScheduler,
    "priority": PriorityScheduler
}

# 动态选择调度器类型
scheduler_cls = schedulers[config.scheduler_type]
# 调用类方法创建实例
scheduler = scheduler_cls.create(config)

在这个例子中，create 类方法在 BaseScheduler 中定义，但 PriorityScheduler 重写了它以提供特定的行为。调用子类的 create 方法时会执行子类的逻辑，然后通过 super().create() 调用父类的实现。

与单例模式结合

vLLM 还将类方法与单例模式结合，确保某些组件只有一个实例：

class TokenizerManager:
    _instances = {}
    
    @classmethod
    def get_tokenizer(cls, model_name: str) -> "Tokenizer":
        """获取或创建对应模型的分词器"""
        if model_name not in cls._instances:
            # 创建新的分词器实例
            cls._instances[model_name] = cls._load_tokenizer(model_name)
        
        # 返回缓存的实例
        return cls._instances[model_name]
    
    @classmethod
    def _load_tokenizer(cls, model_name: str) -> "Tokenizer":
        """加载特定模型的分词器"""
        # 实现分词器加载逻辑
        ...

这种模式确保了资源密集型组件（如分词器）在多次请求时只加载一次，提高了性能和内存效率。

静态方法对比

vLLM 同时使用了静态方法和类方法，它们的区别在于：

class Utility:
    # 静态方法 - 不访问类或实例
    @staticmethod
    def format_tensor(tensor):
        # 纯功能逻辑，不依赖类状态
        return tensor.to(dtype=torch.float16)
    
    # 类方法 - 访问类但不访问实例
    @classmethod
    def create_default(cls):
        # 可以访问类并创建实例
        return cls()

当方法需要访问类本身（如创建实例或访问类变量）时，使用类方法；当方法只是逻辑上与类相关但不需要访问类或实例状态时，使用静态方法。

实际应用案例

在 vLLM 的 LLMEngine 中，类方法用于提供多种创建实例的方式：

class LLMEngine:
    @classmethod
    def from_engine_args(cls, args: EngineArgs) -> "LLMEngine":
        # 从标准参数创建
        ...

    @classmethod
    def from_checkpoint(cls, checkpoint_path: str) -> "LLMEngine":
        # 从检查点恢复
        checkpoint = torch.load(checkpoint_path)
        engine = cls()
        engine.load_state_dict(checkpoint["engine_state"])
        return engine
        
    @classmethod
    def from_pretrained(cls, model_name: str) -> "LLMEngine":
        # 从预训练模型创建
        args = EngineArgs(model=model_name)
        return cls.from_engine_args(args)

这种方式提供了灵活的创建选项，同时保持了接口的一致性和可读性。

总的来说，vLLM 中的类方法和工厂模式是构建灵活且可扩展系统的关键组成部分，它们简化了对象创建，支持多态行为，并提供了清晰的接口供用户使用。

6.6 装饰器高级应用

vLLM 使用多种装饰器增强函数和类的功能，既利用了标准库提供的装饰器，也实现了自定义装饰器：

# 推理模式装饰器
@torch.inference_mode()
def forward(self, *args, **kwargs):
    # 在推理模式下执行，禁用梯度计算
    return self.model(*args, **kwargs)

# 性能分析装饰器
@profile_function
def execute_model(self, batch):
    # 执行并记录性能数据
    return self.model_runner(batch)

# 自定义警告装饰器
@warn_for_unimplemented_methods
class WorkerBase:
    # 抽象基类
    def init_device(self):
        raise NotImplementedError

装饰器是 Python 的强大特性，可以在不修改原始函数定义的情况下添加功能。它们在 vLLM 中被用于多种用途。

torch.inference_mode() 装饰器

torch.inference_mode() 是 PyTorch 提供的性能优化装饰器，类似于 torch.no_grad()，但进一步优化：

@torch.inference_mode()
def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:
    """捕获 CUDA 图以加速推理"""
    # 在此函数中执行的所有 PyTorch 操作都在推理模式下
    # 这会禁用梯度计算，减少内存使用，提高速度
    ...

在推理模式下，PyTorch 会：

禁用梯度跟踪和计算
跳过不必要的内存分配
启用其他优化

对于 vLLM 这种专注于推理的框架来说，这种优化至关重要。

自定义性能分析装饰器

vLLM 实现了自定义装饰器来分析关键函数的性能：

def profile_function(func):
    """记录函数执行时间的装饰器"""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        elapsed_time = time.time() - start_time
        
        # 记录性能数据
        function_name = func.__name__
        logger.debug(f"{function_name} took {elapsed_time:.6f} seconds")
        
        # 更新性能指标
        metrics = get_metrics_singleton()
        metrics.update(function_name, elapsed_time)
        
        return result
    return wrapper

这个装饰器用于包装关键函数，记录它们的执行时间，帮助开发人员识别性能瓶颈。@functools.wraps(func) 确保包装函数保留原始函数的元数据（如名称和文档字符串）。

方法检查装饰器

vLLM 使用装饰器来验证子类是否正确实现了抽象方法：

def warn_for_unimplemented_methods(cls):
    """装饰抽象基类，警告未实现的方法"""
    original_init = cls.__init__
    
    @functools.wraps(original_init)
    def new_init(self, *args, **kwargs):
        # 调用原始 __init__
        original_init(self, *args, **kwargs)
        
        # 检查是否实现了所有抽象方法
        for name, method in inspect.getmembers(cls, inspect.isfunction):
            if getattr(method, "_is_abstract", False):
                # 获取实际实例的方法
                instance_method = getattr(self.__class__, name)
                # 检查是否仍然是抽象方法
                if getattr(instance_method, "_is_abstract", False):
                    logger.warning(
                        f"Class {self.__class__.__name__} does not implement "
                        f"abstract method {name} from {cls.__name__}"
                    )
    
    # 替换 __init__ 方法
    cls.__init__ = new_init
    return cls

这个类装饰器修改了目标类的 __init__ 方法，在初始化时检查子类是否实现了所有标记为抽象的方法，并发出警告。

上下文管理器装饰器

vLLM 使用 contextlib.contextmanager 装饰器创建上下文管理器：

@contextlib.contextmanager
def graph_capture(device: torch.device):
    """CUDA 图捕获上下文管理器"""
    # 设置环境
    torch.cuda.synchronize(device)
    stream = torch.cuda.Stream(device)
    stream.synchronize()
    
    # 获取初始内存状态
    start_free_gpu_memory = torch.cuda.mem_get_info(device)[0]
    
    with torch.cuda.stream(stream):
        try:
            # 进入上下文
            yield {"stream": stream}
        finally:
            # 退出上下文
            stream.synchronize()
            torch.cuda.synchronize(device)
    
    # 记录内存使用
    end_free_gpu_memory = torch.cuda.mem_get_info(device)[0]
    memory_used = start_free_gpu_memory - end_free_gpu_memory

这个装饰器将函数转换为上下文管理器，可以通过 with 语句使用，确保资源的正确分配和释放。

缓存装饰器

vLLM 实现了结果缓存装饰器，用于优化重复计算：

def cache_result(max_size: int = 128):
    """缓存函数结果的装饰器"""
    cache = {}
    order = []
    
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # 创建缓存键
            key = (args, frozenset(kwargs.items()))
            
            # 检查缓存
            if key in cache:
                # 缓存命中，更新使用顺序
                order.remove(key)
                order.append(key)
                return cache[key]
            
            # 缓存未命中，调用原始函数
            result = func(*args, **kwargs)
            
            # 更新缓存
            cache[key] = result
            order.append(key)
            
            # 如果缓存过大，删除最老的项
            if len(cache) > max_size:
                oldest_key = order.pop(0)
                del cache[oldest_key]
                
            return result
        return wrapper
    return decorator

这个装饰器实现了一个简单的 LRU 缓存，可以避免重复计算昂贵的函数调用，特别是对于分词等操作。

组合装饰器

vLLM 经常将多个装饰器组合使用，叠加它们的功能：

# 组合多个装饰器
@torch.inference_mode()
@profile_function
@retry_on_cuda_oom
def forward(self, *args, **kwargs):
    # 函数体
    ...

装饰器的应用顺序是从下到上的，即：

首先应用 @retry_on_cuda_oom
然后应用 @profile_function
最后应用 @torch.inference_mode()

这种组合可以同时提供多种功能增强，如性能优化、错误恢复和执行分析。

总的来说，装饰器在 vLLM 中扮演着重要角色，提供了代码重用、关注点分离和功能增强的优雅方式。它们使得代码更加简洁、可维护，并便于实现横切关注点，如性能分析、错误处理和资源管理。

6.7 上下文管理与多重上下文

vLLM 使用上下文管理器进行资源管理和状态控制，一个典型示例是 CUDA 图捕获过程：

def capture_model(self, kv_caches):
    # 多重嵌套上下文管理器
    with self.attn_state.graph_capture(max_batch_size), \
         graph_capture(self.device) as graph_capture_context, \
         set_forward_context(attn_metadata, self.config, virtual_engine):
        # 执行图捕获
        graph_runner.capture(**capture_inputs)

上下文管理器（通过 with 语句使用）提供了一种优雅的方式来管理资源和状态，确保在代码块执行前后正确地设置和清理环境。

多重嵌套上下文

vLLM 经常使用多重嵌套的上下文管理器来处理复杂的环境设置：

def _execute_with_graphs(self, *args, **kwargs):
    """使用 CUDA 图执行的函数"""
    # 1. 设置流和事件
    with torch.cuda.stream(self.capture_stream):
        # 2. 设置记录状态
        with cuda_profiler.record("cuda_graph_execution"):
            # 3. 设置环境标记
            with set_sync_free_context():
                # 执行图
                self.graph.replay()
                
    # 等待图执行完成
    self.capture_stream.synchronize()

这种嵌套模式确保了正确的资源管理和状态设置顺序，每个上下文都有其特定的职责。

资源管理上下文

vLLM 使用上下文管理器来处理资源分配和释放：

@contextlib.contextmanager
def cuda_memory_context():
    """管理临时 CUDA 内存分配的上下文"""
    # 记录初始内存状态
    initial_allocated = torch.cuda.memory_allocated()
    
    try:
        # 提供上下文
        yield
    finally:
        # 主动释放未使用的缓存内存
        torch.cuda.empty_cache()
        
        # 验证内存是否正确释放
        final_allocated = torch.cuda.memory_allocated()
        if final_allocated > initial_allocated:
            logger.warning(
                f"Memory leak detected: {final_allocated - initial_allocated} bytes"
            )

这个上下文管理器跟踪 CUDA 内存分配，确保在上下文退出时释放未使用的内存，并警告可能的内存泄漏。

状态设置上下文

vLLM 使用上下文管理器来临时修改系统状态：

@contextlib.contextmanager
def set_sync_free_context():
    """设置同步不受限的环境"""
    # 保存原始状态
    original_flag = torch.backends.cuda.matmul.allow_tf32
    
    # 修改状态
    torch.backends.cuda.matmul.allow_tf32 = True
    
    try:
        # 提供上下文
        yield
    finally:
        # 恢复原始状态
        torch.backends.cuda.matmul.allow_tf32 = original_flag

这个上下文管理器临时启用 TF32 数学运算，提高性能，并在退出时恢复原始设置。

异常处理上下文

vLLM 使用上下文管理器来实现更好的异常处理：

@contextlib.contextmanager
def catch_oom_error():
    """捕获并处理 CUDA 内存不足错误"""
    try:
        yield
    except torch.cuda.OutOfMemoryError as e:
        # 清理缓存
        torch.cuda.empty_cache()
        
        # 记录错误
        logger.error("CUDA out of memory. Attempting recovery.")
        
        # 转换为自定义异常
        raise CudaOutOfMemoryError("GPU memory exceeded") from e

这个上下文管理器捕获 CUDA 内存不足错误，尝试清理内存，并转换为自定义异常以便更好地处理。

自定义上下文管理器类

除了使用 @contextlib.contextmanager 装饰器，vLLM 还实现了自定义上下文管理器类：

class GraphContext:
    """CUDA 图捕获和执行的上下文管理器"""
    
    def __init__(self, device, stream=None):
        self.device = device
        self.stream = stream or torch.cuda.Stream(device)
        self.original_stream = None
    
    def __enter__(self):
        # 保存当前流
        self.original_stream = torch.cuda.current_stream(self.device)
        
        # 同步设备
        torch.cuda.synchronize(self.device)
        
        # 设置新流
        self.stream.synchronize()
        torch.cuda.set_stream(self.stream)
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 同步流
        self.stream.synchronize()
        
        # 恢复原始流
        torch.cuda.set_stream(self.original_stream)
        
        # 同步设备
        torch.cuda.synchronize(self.device)

这个类实现了 __enter__ 和 __exit__ 方法，提供了更复杂的上下文管理功能，特别适合需要精确控制 CUDA 流的场景。

嵌套上下文的工作原理

当使用多重嵌套上下文时，工作流程如下：

按从外到内的顺序调用每个上下文管理器的 __enter__ 方法
执行 with 块中的代码
按从内到外的顺序调用每个上下文管理器的 __exit__ 方法

这确保了资源和状态的正确设置和清理顺序，即使在异常发生时也是如此：

# 嵌套上下文的工作流程
# 1. memory_context.__enter__()
# 2. graph_context.__enter__()
# 3. profile_context.__enter__()
# 4. 执行代码块
# 5. profile_context.__exit__()
# 6. graph_context.__exit__()
# 7. memory_context.__exit__()

with cuda_memory_context():
    with GraphContext(device) as graph:
        with cuda_profiler.record("operation"):
            # 执行代码
            model.forward(inputs)

即使在代码执行期间发生异常，所有上下文管理器的 __exit__ 方法仍会被调用，确保资源得到正确清理。

实际应用案例

在 vLLM 的 capture_model 方法中，多重上下文管理器确保了 CUDA 图捕获的正确环境：

@torch.inference_mode()
def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:
    # ...准备工作...
    
    # 多重上下文管理器
    with self.attn_state.graph_capture(max_batch_size), \
         graph_capture(self.device) as graph_capture_context, \
         set_forward_context(attn_metadata, self.config, virtual_engine):
        
        for batch_size in cudagraph_capture_sizes:
            # ...准备捕获输入...
            
            # 创建并捕获图
            graph_runner = CUDAGraphRunner(device=self.device)
            graph_runner.capture(**capture_inputs)
            
            # 保存图
            self.graph_runners[batch_size] = graph_runner

每个上下文管理器负责特定方面：

self.attn_state.graph_capture 准备注意力状态
graph_capture 设置 CUDA 流和同步点
set_forward_context 配置前向传播环境

总的来说，上下文管理器在 vLLM 中扮演着关键角色，提供了优雅的资源管理和状态控制机制，确保了复杂操作（如 CUDA 图捕获）的正确执行环境，同时简化了错误处理和资源清理。这种机制对于管理 GPU 资源和确保系统稳定性特别重要。

7. 性能优化与最佳实践

7.1 示例应用：构建高性能API服务

vLLM可以轻松构建高性能的LLM API服务，以下是一个完整的示例：

from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel, Field
from typing import List, Optional
from vllm import LLM, SamplingParams
import uvicorn
import asyncio

app = FastAPI(title="高性能LLM API服务")

# 初始化模型（在生产环境中应使用更完整的配置）
llm = LLM(
    model="llama-7b",
    tensor_parallel_size=2,  # 使用2个GPU进行张量并行
    gpu_memory_utilization=0.8,
    max_num_seqs=256,  # 批处理大小
)

class GenerationRequest(BaseModel):
    prompt: str
    max_tokens: int = Field(default=128)
    temperature: float = Field(default=0.7)
    top_p: float = Field(default=0.9)
    stream: bool = Field(default=False)

class GenerationResponse(BaseModel):
    text: str
    usage: dict

@app.post("/generate", response_model=GenerationResponse)
async def generate(request: GenerationRequest):
    # 设置采样参数
    params = SamplingParams(
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
    )
    
    # 异步执行生成
    outputs = await asyncio.to_thread(
        llm.generate, [request.prompt], params
    )
    
    # 构建响应
    response = GenerationResponse(
        text=outputs[0].outputs[0].text,
        usage={
            "prompt_tokens": len(outputs[0].prompt_token_ids),
            "completion_tokens": len(outputs[0].outputs[0].token_ids),
            "total_tokens": len(outputs[0].prompt_token_ids) + 
                           len(outputs[0].outputs[0].token_ids)
        }
    )
    
    return response

# 添加健康检查端点
@app.get("/health")
async def health_check():
    return {"status": "healthy"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

这个服务具有以下特点：

高并发处理：利用vLLM的批处理能力处理多个并发请求
资源高效利用：通过张量并行和内存优化最大化GPU利用率
生产级API：提供标准的REST API接口，兼容其他LLM服务形式
性能监控：可以通过添加中间件来监控请求延迟和吞吐量

7.2 性能优化策略

7.2.1 批处理优化

vLLM通过连续批处理（Continuous Batching）技术显著提高了吞吐量：

静态批处理的局限：传统批处理要求所有序列同时启动和结束，长序列会使短序列等待，导致GPU利用率低
连续批处理优势：允许动态添加和移除序列，提高GPU利用率达80-90%

实现方法：

# 优化批处理代码示例
def _prepare_batch(self, scheduled_seq_groups):
    batch = Batch()
    # 动态收集当前活跃序列
    for seq_group in scheduled_seq_groups:
        batch.add_sequence_group(seq_group)
    # 处理不同长度的序列
    return batch.finalize()

7.2.2 内存优化

PagedAttention：将KV缓存组织为固定大小的物理块，通过块表映射实现高效内存管理
量化技术：支持INT8/INT4量化，可将模型内存需求减少50-75%
激活检查点：在推理中重计算某些中间结果，而非全部保存，平衡计算和内存使用
内存使用监控：动态调整批大小和缓存策略，避免OOM错误

7.2.3 计算优化

CUDA图优化：预先捕获和优化GPU操作序列，降低CPU-GPU通信开销
自定义CUDA核心：针对PagedAttention实现了高效的CUDA核心
张量并行：将大型模型分割到多个GPU上执行
异步执行：利用流水线执行和异步通信提高并发性

7.3 部署最佳实践

7.3.1 硬件选择与配置

GPU选择：推荐使用支持CUDA的NVIDIA GPU，如A100、H100系列
内存配置：
- 计算模型内存需求：模型大小 + 批处理KV缓存 + 激活内存
- 对于14B模型，推荐至少40GB显存
CPU配置：CPU核心数影响分词和后处理性能，建议每个GPU配置4-8个CPU核心

7.3.2 系统参数调优

批处理大小：根据模型大小和GPU内存动态调整

# 批处理大小计算示例
max_batch_size = (gpu_memory_gb * 0.9 - model_size_gb) // (
    token_kv_cache_size_gb * max_seq_len)

块大小：KV缓存的物理块大小，默认16，可根据应用场景调整为8-32

预分配缓存：对于固定工作负载，预分配缓存可减少动态分配开销

engine = LLMEngine(
    model="llama-7b",
    block_size=16,
    gpu_memory_utilization=0.9,
    swap_space=4,  # GB
)

7.3.3 分布式部署

单机多GPU：优先使用张量并行，设置tensor_parallel_size参数
多机部署：确保机器间高速网络连接（推荐InfiniBand或100Gb以太网）
负载均衡：使用轮询或最小负载策略分发请求

7.3.4 常见性能问题及解决方案

问题	可能原因	解决方案
GPU利用率低	批处理大小不足	增加`max_num_seqs`或`max_num_batched_tokens`
内存溢出 (OOM)	KV缓存过大	减小批大小或启用缓存压缩
高延迟	CPU瓶颈	增加CPU核心或优化分词逻辑
吞吐量不足	CUDA图未启用	确保`disable_cuda_graphs=False`
输入处理慢	Tokenizer配置	使用快速分词器或预处理批量输入

7.4 常见问题与解决方案

问题	可能原因	解决方案
GPU利用率低	批处理大小不足	增加`max_num_seqs`或`max_num_batched_tokens`
内存溢出 (OOM)	KV缓存过大	减小批大小或启用缓存压缩
高延迟	CPU瓶颈	增加CPU核心或优化分词逻辑
吞吐量不足	CUDA图未启用	确保`disable_cuda_graphs=False`
输入处理慢	Tokenizer配置	使用快速分词器或预处理批量输入

8. 总结与展望

8.1 技术创新点总结

vLLM作为高性能大语言模型推理框架，其核心创新点包括：

技术创新	核心价值	实现方式
PagedAttention	解决KV缓存内存碎片问题	将连续logical序列映射到不连续物理内存块
连续批处理	提高GPU利用率	动态调度机制允许序列独立进出批处理队列
CUDA图捕获	减少CPU-GPU通信开销	预编译和优化GPU操作序列
分布式推理	支持超大模型部署	实现张量并行和流水线并行
动态加载卸载	高效内存管理	通过Python反射机制实现按需加载

vLLM通过上述创新，在保持API简洁性的同时，实现了高达4倍的性能提升，以及更高效的内存使用。其设计理念体现了现代软件工程的最佳实践，包括：

模块化设计：清晰的组件分离使得系统易于扩展和维护
抽象接口：通过抽象接口隔离实现细节，支持多种硬件后端
高级语言特性：巧妙利用Python的动态特性实现灵活配置和加载
并发与并行：多层次的并行策略最大化硬件利用率

8.2 与其他框架对比

以下是vLLM与其他主流LLM推理框架的对比：

特性	vLLM	HuggingFace TGI	DeepSpeed-Inference	TensorRT-LLM
连续批处理	✅	✅	❌	✅
PagedAttention	✅	❌	❌	❌
张量并行	✅	✅	✅	✅
流水线并行	✅	❌	✅	✅
量化支持	INT8/INT4	INT8	INT8/INT4	INT8/INT4/FP8
自定义模型支持	中等	高	中等	低
开发难度	低	低	高	高
性能表现	优秀	良好	优秀	极佳
部署复杂度	低	低	高	高
社区活跃度	高	高	中	中

vLLM在易用性和性能之间取得了良好平衡，特别适合快速部署和优化大语言模型服务。

8.3 未来发展趋势

根据vLLM的发展轨迹和大语言模型服务的行业需求，可以预见以下发展方向：

8.3.1 技术演进方向

更多硬件后端支持：扩展到AMD/Intel GPU和专用AI加速器
混合精度推理：根据不同层的敏感性动态调整精度
分布式推理增强：改进跨数据中心的模型部署效率
自适应批处理：基于负载动态调整批处理策略
内存优化：KV缓存压缩和选择性保存
推理时微调：支持推理过程中的参数高效微调

8.3.2 应用场景扩展

边缘设备部署：优化小型设备上的模型运行效率
多模态支持：扩展框架以支持图像、音频等多模态输入
专用场景优化：为RAG、Agent等应用场景提供专门优化
服务化增强：提供更完善的监控、扩缩容和灾备功能

8.4 实践建议

对于开发者和研究人员，建议从以下几个方面入手深入理解和应用vLLM：

从小模型开始实践：先用较小模型熟悉API和工作流程
理解内存分配机制：掌握PagedAttention的核心原理
性能分析与调优：学习使用profiler工具分析性能瓶颈
扩展自定义功能：基于vLLM架构开发特定应用场景的优化
关注社区动态：vLLM发展迅速，定期关注新特性和优化

通过深入理解vLLM的设计理念和技术实现，开发者可以更好地应用和扩展这一框架，为大语言模型的落地应用提供强有力的技术支持。

附录

A. 关键源码目录结构

vllm/
├── config.py           # 配置类定义
├── engine/             # 核心引擎
│   ├── arg_utils.py    # 参数处理
│   ├── async_llm_engine.py # 异步引擎
│   └── llm_engine.py   # LLM引擎主类
├── entrypoints/        # 入口点
│   ├── api_server.py   # API服务器
│   ├── llm.py          # 命令行接口
│   └── openai.py       # OpenAI兼容API
├── executor/           # 执行器
│   ├── executor_base.py# 执行器基类
│   └── schedulers/     # 调度器
├── model_executor/     # 模型执行器
│   ├── layers/         # 模型层实现
│   └── models/         # 模型定义
├── sampling/           # 采样策略
│   ├── sampling_params.py # 采样参数
│   └── samplers.py     # 采样器实现
├── sequence.py         # 序列管理
├── worker/             # 工作节点
│   ├── model_runner.py # 模型运行器
│   └── worker.py       # 工作节点实现
└── utils.py            # 工具函数

B. 参考资源

Mcp 智能聊天助手示例项目

Wed, 19 Mar 2025 00:00:00 +0000

MCP 智能聊天助手示例项目

项目地址

基于Model Context Protocol的智能对话系统示例

项目简介

本项目是基于MCP (Model Context Protocol)框架的智能聊天助手示例，展示了如何使用大语言模型(LLM)与外部数据源和工具进行集成，实现跨模态、多能力的智能对话系统。

MCP框架介绍

MCP (Model Context Protocol)是由Anthropic开源的开放协议，为LLM应用和外部数据源及工具之间的无缝集成提供标准化方式。可以将MCP理解为AI应用的”USB-C接口”，就像USB-C为设备提供了与各种外设和配件连接的标准方式一样，MCP为AI模型提供了与不同数据源和工具连接的标准方式。

MCP核心价值：连接用户、大模型与专业服务

MCP本质上是一个标准化的连接协议，它作为关键的桥梁，将三个核心组件优雅地连接在一起：

用户 - 通过自然语言表达需求的人或者其他的输入者
大模型 - 处理理解和生成的AI系统
专业服务 - 提供特定领域能力的工具和数据源

作为连接器，MCP实现了以下关键价值：

降低用户门槛：允许用户用自然语言表达复杂需求，无需了解底层技术细节
扩展模型能力：让大模型能够访问实时外部数据和专业功能，克服了模型知识局限和”幻觉”问题
激活专业服务：使现有专业服务能够被大模型智能调用，无需彻底重构现有系统
标准化交互：通过统一协议减少集成复杂性，降低开发和维护成本

这种标准协议的设计使得整个AI应用生态系统更加模块化和可扩展，各组件可以专注于自己的强项，同时通过标准接口协同工作，为用户提供集成化的智能体验。

为什么选择MCP?

MCP帮助你在LLM之上构建智能代理和复杂工作流。大语言模型经常需要与数据和工具集成，而MCP提供：

预构建集成列表，可直接连接到你的LLM
在不同LLM提供商和供应商之间切换的灵活性
在你的基础设施内保护数据的最佳实践

MCP架构

在核心上，MCP遵循客户端-服务器架构，其中主机应用程序可以连接到多个服务器：

MCP主机：如Claude Desktop、IDE或AI工具等希望通过MCP访问数据的程序
MCP客户端：维护与服务器1:1连接的协议客户端
MCP服务器：通过标准化的Model Context Protocol暴露特定功能的轻量级程序
本地数据源：MCP服务器可以安全访问的计算机文件、数据库和服务
远程服务：通过互联网可用的外部系统（例如API），MCP服务器可以连接到这些系统

MCP工作流程

MCP的实际工作流程遵循以下步骤：

服务准备：首先，你需要一个或多个遵循MCP标准的MCP服务器，每个服务器为特定领域暴露一个功能接口（工具）和数据资源。这些服务器可以是本地的或远程的。
用户应用：用户通过集成了MCP客户端的应用程序（如聊天界面、IDE插件等）发送请求，MCP客户端负责与MCP服务器通信。
请求处理：
- 用户发送自然语言请求（例如”查询明天北京的天气”）
- 应用程序收集所有已注册的MCP服务器及其工具信息
- 应用程序调用大语言模型分析用户意图，并将其映射到适当的MCP服务和工具
- 模型返回结构化调用信息（服务名称、工具名称、参数等）
服务调用：
- 应用程序根据模型的决定，调用指定MCP服务器的相关工具
- MCP服务器处理请求并返回结构化数据
结果处理：
- 应用程序从服务器接收专业数据
- 再次调用模型将专业数据转换为用户友好的自然语言回复
- 向用户呈现最终回复

下图展示了MCP的工作流程：

┌───────────────┐        ┌───────────────┐        ┌───────────────┐
│               │        │               │        │               │
│     用户      │───────▶│  应用程序     │◀──────▶│   大语言模型  │
│               │        │ (MCP客户端)   │        │               │
└───────────────┘        └───────▲───────┘        └───────────────┘
                                 │
                                 │
                         ┌───────▼───────┐
                         │  MCP服务注册  │
                         │    和发现     │
                         └───────┬───────┘
                                 │
                 ┌───────────────┼───────────────┐
                 │               │               │
         ┌───────▼───────┐┌──────▼────────┐┌─────▼─────────┐
         │               ││               ││                │
         │ MCP服务器 A   ││  MCP服务器 B  ││  MCP服务器 C   │
         │ (天气服务)    ││  (知识库)     ││  (数据分析)    │
         │               ││               ││                │
         └───────────────┘└───────────────┘└────────────────┘

这个工作流程展示了MCP的核心价值：允许大语言模型通过标准化协议无缝利用各种专业服务，同时保持一致的用户体验。用户只需通过自然语言表达需求，无需了解底层技术细节或服务架构。

MCP多服务器协调机制

MCP框架能够优雅地协调多个服务器，实现复杂功能集成。具体协调机制包括：

服务发现和注册：
- MCP主机在启动时扫描并发现本地或网络中的MCP服务器
- 每个服务器在启动时向主机注册其能力（工具、资源类型等）
- 主机维护服务器能力映射表，用于路由请求
能力描述和工具列表：
- 每个MCP服务器发布其支持的工具列表、参数规范和资源类型
- 主机可以通过list_tools和list_resources等API获取每个服务器的能力
- 模型根据这些能力决定调用哪个服务器的哪个工具
上下文管理：
- MCP维护请求上下文，以确保多服务器调用之间的信息一致性
- 上下文可以包括会话信息、用户偏好、身份验证令牌等
- 服务器可以读取上下文以提供个性化响应
请求路由：
- 当应用程序或模型需要特定服务时，MCP客户端负责：
  - 识别哪个服务器可以处理请求
  - 将请求格式化为标准的MCP协议消息
  - 与目标服务器建立连接
  - 传输请求并等待响应
并行处理：
- MCP支持同时向多个服务器发送请求
- 可以并行获取不同数据源的信息
- 可以将多个服务器的结果聚合到一个响应中
错误处理和回退：
- 当特定服务器不可用时，MCP可以自动尝试备用服务器
- 错误报告机制帮助应用程序诊断问题
- 支持为服务中断定义回退策略
版本兼容性：
- MCP协议设计支持不同版本的服务器共存
- 客户端可以处理版本差异，以确保向后兼容性
- 新服务器可以提供增强功能，而无需更新整个环境

在实践中，例如，当用户问”分析上个月的销售数据并与天气比较”时，MCP将：

识别需要销售数据分析和天气查询服务
调用数据分析MCP服务器获取销售统计信息
调用天气MCP服务器获取历史天气数据
将两种数据组合并返回给模型
模型根据完整上下文生成综合分析

这种协调机制使MCP能够无缝结合多个专业服务器的能力，提供集成化的体验，而每个服务器只专注于自己的专业领域，大大提高了系统的模块化和可扩展性。

MCP框架的核心思想是为大语言模型应用提供结构化的上下文信息，包括：

资源：提供给模型的数据，如文件内容、API响应等
工具：模型可以调用的函数，用于执行操作和获取信息
提示：可重用的交互模板，用于引导模型响应

MCP实际应用场景

让我们通过几个生动的例子来看看MCP如何解决实际问题：

例1: 智能代码助手

问题：传统代码助手缺乏对代码仓库的深入理解，无法访问本地代码结构和版本控制信息。

MCP解决方案：

开发者安装支持MCP的IDE插件
IDE创建MCP服务器，暴露本地文件系统、Git历史记录和代码索引
当开发者通过MCP向AI提问”这个bug可能的原因是什么？”时，AI可以：
- 读取当前打开文件的内容
- 查询错误日志
- 检索相关代码模块
- 分析Git提交历史以找到最近的更改
- 提供具体的、基于上下文的解决方案

效果：开发者获得基于项目实际情况的具体建议，而不是泛泛而谈的答案。

例2: 私密知识库问答

问题：企业希望员工能够询问内部文档和数据，但担心隐私泄露。

MCP解决方案：

企业在自己的基础设施上部署MCP服务器，连接内部知识库
员工在聊天界面中问”2023年第三季度我们的销售策略是什么？”
MCP服务器在本地：
- 检索相关内部文档
- 提取所需信息
- 将内容发送给模型进行总结
- 员工收到基于内部数据的答案

效果：敏感数据留在企业内部，但员工仍然获得智能问答体验。

例3: 个人数据助手

问题：用户希望AI帮助管理个人数据，但担心隐私问题。

MCP解决方案：

用户在自己的设备上安装MCP应用程序
应用程序创建本地服务器，连接个人照片库、日历和文件
用户问”帮我找到去年夏天在海滩拍的照片，并创建一个相册”
MCP服务器在本地处理：
- 在照片库中搜索相关图像
- 仅将照片缩略图发送给模型进行分析
- 根据模型建议在本地设备上创建相册

效果：用户获得AI辅助的个人数据管理能力。

例4: 多源研究助手

问题：研究人员需要处理多种不同格式的数据并切换工具耗时。

MCP解决方案：

研究人员部署多个MCP服务器，分别连接：
- 科学文献数据库
- 数据可视化工具
- 统计分析软件
- 笔记和引用管理系统
研究员可以问”分析这组基因数据与近期论文中的发现有何关联”
MCP协调多个服务器：
- 检索并分析实验数据
- 搜索最新相关论文
- 运行统计比较
- 生成可视化图表
- 整合为一份综合报告

效果：研究人员通过一个统一界面完成跨数据源、跨工具的复杂工作流。

这些例子说明了MCP如何解决数据集成、隐私保护和工具互操作性等关键挑战，使AI应用能够安全地访问和处理各种数据源，同时提供丰富的上下文，让大语言模型发挥最大价值。

项目架构

本项目采用了模块化设计，主要包含以下组件：

聊天服务(Chat App)：核心服务组件，负责
- 用户请求处理
- 大模型API调用
- 用户意图解析
- MCP服务调用协调
- 响应格式化
MCP管理器：负责
- MCP服务发现和注册
- 工具列表管理
- 服务调用路由
天气MCP服务：提供天气信息查询功能
- 天气数据获取
- 温度范围查询
- 降水预测
- 天气图表生成
大语言模型服务(vLLM)：提供大模型推理能力
- 模型托管和优化推理
- 用户意图解析和服务选择
- 结构化数据转自然语言描述
- OpenAI兼容API接口

架构图示：

┌─────────────────────────┐           ┌──────────────────────┐
│                         │           │                      │
│      用户/客户端        │◀─HTTP────▶│    mcp_chat_app/     │
│                         │           │    chat_app.py       │
└─────────────────────────┘           │    (Flask应用)       │
                                      └──────────┬───────────┘
                                                 │
                                                 │
                  ┌──────────────────────────────┼──────────────────────────────┐
                  │                              │                              │
     ┌────────────▼─────────────┐   ┌────────────▼─────────────┐   ┌────────────▼─────────────┐
     │  mcp_chat_app/           │   │   mcp_chat_app/          │   │     vllm_server/         │
     │  mcp_manager.py          │   │   mock_mcp_manager.py    │   │     run.py               │
     │  (MCP服务发现与管理)     │   │   (测试用模拟MCP管理器)  │   │   (大模型推理服务        │
     └────────────┬─────────────┘   └──────────────────────────┘   │    可替换为远程API)      │
                  │                                                 └──────────┬───────────────┘
                  │ MCP协议                                                    │
                  │                                                            │
     ┌────────────▼─────────────┐                                  ┌───────────▼───────────────┐
     │  weather_mcp_server/     │◀─────────────────────────────────┤ 1. 意图解析              │
     │  weather_server.py       │                                  │ 2. 参数提取和验证         │
     │  (天气服务MCP实现)       │─────────────────────────────────▶│ 3. 结果格式化和润色      │
     │                          │                                  └───────────────────────────┘
     └──────────────────────────┘                                             ▲
             │        │                                                       │
             │        │                                                       │
┌────────────▼──┐ ┌───▼────────────┐           ┌───────────────────────────┐ │
│ utils.py      │ │ chart_generator.│           │       双向交互流程        │ │
│ (数据处理)    │ │ py (图表生成)   │           │                           │ │
└───────────────┘ └────────────────┘           │ 1. 用户查询意图理解        │ │
                                               │ 2. MCP工具选择             │ │
                                               │ 3. 参数解析和转换          │ │
                                               │ 4. 结果数据自然语言化      │ │
                                               │ 5. 多模态内容处理          │ │
                                               └───────────────────────────┘ │
                                                                             │
                                                       chat_app.py ──────────┘

大模型服务关键交互

大语言模型服务(vLLM)在整个系统中扮演着”大脑”角色，与其他组件有三个关键交互点：

与聊天服务(chat_app.py)的交互：
- 接收用户原始查询，分析用户意图
- 从可用工具列表中选择最合适的MCP服务和工具
- 提取和格式化调用参数
- 将MCP服务返回的结构化数据转化为自然语言回复
- 处理错误情况并提供友好响应
与MCP服务的间接交互：
- 通过聊天服务作为中介，大模型决定调用哪个MCP服务
- 根据MCP服务的工具文档信息，正确构造调用参数
- 理解MCP服务返回的数据结构，进行信息提取和整合
模型能力应用：
- 语义理解：将自然语言转化为结构化意图
- 上下文记忆：维持多轮对话的连贯性
- 格式转换：在自然语言和结构化数据间互相转换
- 多模态处理：支持文本、图表、结构化数据等混合内容

大模型服务采用vLLM加速推理，提供OpenAI兼容的API接口，使系统能够轻松替换底层模型，同时保持接口一致性。该服务不仅提供基础推理能力，还承担着整个系统的决策和协调角色，是连接用户需求与专业服务的关键桥梁。

值得注意的是，本架构中的vLLM服务器组件采用了模块化设计，可以灵活替换为：

本地部署模式：如当前的vllm_server部署，适合开发测试或对数据安全有较高要求的场景
云端API模式：可直接替换为OpenAI、Anthropic Claude、QianWen千问等商业API服务
自托管远程模式：可部署在专用GPU服务器或云端GPU实例上，通过API连接
混合模式：不同功能使用不同来源的模型服务，如意图理解使用轻量模型，内容生成使用更强大的模型

系统设计确保了无论使用哪种模型服务部署方式，聊天应用的核心功能和与MCP服务的交互方式都保持一致，只需修改配置即可切换，无需更改应用代码架构。这种灵活性使开发者可以根据实际需求、预算和计算资源选择最适合的部署策略。

系统工作流程：

用户发送HTTP请求到聊天服务(chat_app.py)
聊天服务使用vLLM大模型服务分析用户意图
根据分析结果，通过mcp_manager调用天气MCP服务的相关工具
天气服务处理请求，可能使用chart_generator生成图表
聊天服务接收MCP服务返回的专业数据
再次调用大模型服务将数据转化为自然语言回复
将最终回复返回给用户

用户使用方式

通过调用api，在api中用户可以通过自然语言查询天气信息：

“北京今天天气怎么样？”
“上海未来三天的温度变化趋势”
“广州今天会下雨吗？”
“给我看看杭州的降水量图表”

系统会调用天气服务获取相关数据，并通过大模型将数据转化为易于理解的回复。

如何用cursor开发自己的MCP服务(以当前项目，天气查询服务为例)

创建一个新的文件夹
拷贝mcp的readme以及python sdk readme到当前的项目根目录以告诉cursor，mcp是什么以及怎么用

生成mcpserver

使用cursor agent开始第一个项目生成

请根据REAME中的MCP框架创建一个项目，这个项目是一个mcpserver，实现了mcpserver的所有功能，作天气查询

这样就会生成 weather_mcp_server 项目

运行mcp

为整个项目创建一个 venv

# 回到项目根目录
uv venv -p python3.12 .venv
source .venv/bin/activate
# 安装mcp
pip install mcp 'mcp[cli]'
# 安装依赖
pip install -r weather_mcp_server/requirements.txt
# 运行mcpserver
python weather_mcp_server/run.py --dev
# 或者使用mcp命令
mcp dev weather_mcp_server/weather_server.py
# 输出如下
Starting MCP inspector...
Proxy server listening on port 3000

🔍 MCP Inspector is up and running at http://localhost:5173 🚀

这样我们就可以在http://localhost:5173 这个mcp代理中访问mcpserver了

在claude和cursor中都可以使用，例如在cursor中使用

进行添加会帮我们创建一个文件，文件内容需要根据刚才创建的mcpserver来填写

{
 "mcpServers": {
   "server-name": {
  "command": "python",
  "args": ["/yourpath/weather_mcp_server/weather_server.py"]
 }
  }
}

注册上来之后在cursor agent中可以通过cursor使用mcpserver

生成用户程序

现在mcpserver使用依赖于cursor，mcp设计本身事为了作为大模型与应用的连接器的，所以我们能不能写个应用来连接大模型和应用呢，现在应用已经有了，需要

有一个大模型(本地启动的无论是deepseek还是、llama、qianwen都表现比较差，用阿里云百联的qwen-plus)
- 可以使用项目中的vllm代码本地启动一个
- 使用在线大模型
有一个面向用户的程序，用户发送请求，例如”北京天气如何”，此程序通过大模型分析用户语义并和mcpserver暴露的服务做匹配，在调用对应服务获取返回，大模型加工之后返回给用户,可参考mcp_simple_chatbot
- 需要有mcpserver注册功能
- 解决上述业务流程

将上面这段话发送给cursor agent，便能生成项目 mcp_chat_app，经过调试基本可用

curl -X POST "http://localhost:8081/chat/test_user" -H "Content-Type: application/json" -d '{"message":  "明天29以上，30以下的城市有哪些？", "use_tools": true}'
{"response":"根据天气预报，明天温度在29°C到30°C之间的城市有以下两个：\n\n1. **广州**：预计最高温度为30°C，天气以高温为主。\n2. **深圳**：预计温度为29.5°C，天气多云。\n\n如果你计划去这两个城市之一，记得根据天气情况做好准备哦！ 😊"}%     

curl -X POST "http://localhost:8081/chat/test_user" -H "Content-Type: application/json" -d '{"message":  "西安今天天气如何？", "use_tools": true}'
{"response":"您好！关于西安今天的天气，目前MCP服务没有直接提供具体信息，但它给出了未来7天内部分城市的天气概况。从这些数据中，我们可以看到：\n\n- **最热的城市**是广州和深圳，分别达到30.0°C（高温）和29.5°C（多云）。  \n- **有降雨的城市**是成都，预报为小雨。\n\n虽然西安的天气未在上述列表中提及，但通常情况下，西安的天气以干燥和晴朗为主，建议您参考本地天气应用或稍后再查询以获取更准确的信息哦！如果需要进一步帮助，请告诉我！ 😊"}%      

项目扩展

本项目可以通过以下方式进行扩展:

添加更多MCP服务: 开发新的专业服务，如文档处理、代码生成等
增强聊天体验: 添加上下文管理、多轮对话支持
提升大模型能力: 接入更强大的模型、优化提示词工程
开发Web/移动前端: 构建用户友好的界面
添加身份验证: 实现用户管理和权限控制

MCP当前生态

MCP作为一项相对新兴的技术标准，其生态系统目前仍处于发展初期：

成熟应用:
- Claude桌面应用是目前集成MCP最完善的应用，提供了完整的MCP服务发现、管理和调用功能
- Cursor IDE已实现基础的MCP支持，允许开发者在编码过程中利用本地上下文
生态状况:
- 当前支持MCP的应用和工具相对有限，主要集中在开发工具和专业领域
- 官方SDK支持主流编程语言(Python、TypeScript、Java、Kotlin)，但社区贡献的MCP服务库尚在起步阶段
- 大多数MCP服务需要开发者自行构建和部署
发展趋势:
- 随着Anthropic持续推动标准发展和生态建设，更多应用预计将添加MCP支持
- 厂商中立的特性使得MCP有潜力成为连接LLM和专业服务的通用标准
- 开发者社区对私有部署和数据安全的需求正在推动MCP在企业环境中的应用

对于开发者而言，现阶段是参与MCP生态建设的有利时机，可以通过构建专业领域的MCP服务来填补生态空白，也可以将现有工具和服务通过MCP协议进行包装，以便与大模型应用集成。

贡献指南

欢迎为项目做出贡献! 您可以通过以下方式参与:

提交问题和功能请求
提交代码改进
贡献新的MCP服务
改进文档

许可证

Apache-2.0 License

使用3090显卡部署Wan2.1生成视频

Fri, 28 Feb 2025 00:00:00 +0000

使用3090显卡部署Wan2.1生成视频

使用3090显卡部署Wan2.1生成视频

1. 环境说明

名称	规格	备注
内存VRAM	22GB	22G不够用，需要更大的内存，如果没有的话可以用swap内存代替
显卡	3090	NVIDIA
显存	24GB
CUDA	12.5	CUDA >= 11.7 否则会失败

2. 模型下载

使用 hfs.sh脚本下载模型

./hfd.sh Wan-AI/Wan2.1-T2V-1.3B --tool aria2c -x 10 --hf_token xxxxxxx --hf_username xxxxxxx

3. 克隆仓库

git clone git@github.com:Wan-Video/Wan2.1.git

4. 安装依赖

Note: CUDA >= 11.7 torch >= 2.4.0

cd Wan2.1
pip install -r requirements.txt

5. 生成视频

5.1. 使用generate脚本生成

Note: 一定要配置 –offload_model True –t5_cpu 否则GPU内存不够会在保存视频时OOM

 python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --offload_model True --t5_cpu --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
[2025-02-28 08:33:08,062] INFO: Generation job args: Namespace(task='t2v-1.3B', size='832*480', frame_num=81, ckpt_dir='./Wan2.1-T2V-1.3B', offload_model=True, ulysses_size=1, ring_size=1, t5_fsdp=False, t5_cpu=True, dit_fsdp=False, save_file=None, prompt='Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.', use_prompt_extend=False, prompt_extend_method='local_qwen', prompt_extend_model=None, prompt_extend_target_lang='ch', base_seed=6930324173022001627, image=None, sample_solver='unipc', sample_steps=50, sample_shift=8.0, sample_guide_scale=6.0)
[2025-02-28 08:33:08,063] INFO: Generation model config: {'__name__': 'Config: Wan T2V 1.3B', 't5_model': 'umt5_xxl', 't5_dtype': torch.bfloat16, 'text_len': 512, 'param_dtype': torch.bfloat16, 'num_train_timesteps': 1000, 'sample_fps': 16, 'sample_neg_prompt': '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走', 't5_checkpoint': 'models_t5_umt5-xxl-enc-bf16.pth', 't5_tokenizer': 'google/umt5-xxl', 'vae_checkpoint': 'Wan2.1_VAE.pth', 'vae_stride': (4, 8, 8), 'patch_size': (1, 2, 2), 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'num_heads': 12, 'num_layers': 30, 'window_size': (-1, -1), 'qk_norm': True, 'cross_attn_norm': True, 'eps': 1e-06}
[2025-02-28 08:33:08,063] INFO: Input prompt: Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.
[2025-02-28 08:33:08,065] INFO: Creating WanT2V pipeline.
[2025-02-28 08:35:19,682] INFO: loading ./Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth
[2025-02-28 08:37:19,359] INFO: loading ./Wan2.1-T2V-1.3B/Wan2.1_VAE.pth
[2025-02-28 08:37:24,292] INFO: Creating WanModel from ./Wan2.1-T2V-1.3B
[2025-02-28 08:38:17,205] INFO: Generating video ...
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [08:15<00:00,  9.91s/it]
[2025-02-28 08:47:29,986] INFO: Saving generated video to t2v-1.3B_832*480_1_1_Two_anthropomorphic_cats_in_comfy_boxing_gear_and__20250228_084729.mp4
[2025-02-28 08:47:32,176] INFO: Finished.

视频生成完成, 可以在当前目录下查看生成的视频

t2v-1.3B_832*480_1_1_Two_anthropomorphic_cats_in_comfy_boxing_gear_and__20250228_084729.mp4

您的浏览器不支持 video 标签。

5.2. 使用gradio启动UI界面生成

5.2.1. 启动gradio服务

这里需要一个prompt扩展的API，可以使用dashscope的API，也可以使用自己的API，这里使用dashscope的API https://dashscope.console.aliyun.com/overview

DASH_API_KEY=xxxx  python gradio/t2v_1.3B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir ./Wan2.1-T2V-1.3B/
Step1: Init prompt_expander...done
Step2: Init 1.3B t2v model...done
* Running on local URL:  http://0.0.0.0:7860

To create a public link, set `share=True` in `launch()`.
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [08:14<00:00,  9.88s/it]

5.2.2. 访问UI，输入prompt，点击生成视频

在白雪皑皑的梅里雪山上空，数架战斗机编队飞行，战斗机在空中变换各种队形和特技动作，背景是晴朗的蓝天和壮丽的山峰，画面充满动感和震撼力。

Over the snow-covered Meili Snow Mountains, multiple fighter jets fly in formation, performing various aerial maneuvers and stunts against a backdrop of clear blue skies and majestic peaks. The scene is dynamic and awe-inspiring.

您的浏览器不支持 video 标签。

Deploying Wan2.1 for Video Generation with 3090 GPU

Fri, 28 Feb 2025 00:00:00 +0000

Deploying Wan2.1 for Video Generation with 3090 GPU

Deploying Wan2.1 for Video Generation with 3090 GPU

1. Environment Requirements

Name	Specification	Notes
Memory VRAM	22GB	22GB is not enough, needs more memory. If not available, swap memory can be used as an alternative
GPU	3090	NVIDIA
VRAM	24GB
CUDA	12.5	CUDA >= 11.7 required, otherwise it will fail

2. Model Download

Use hfs.sh script to download the model

./hfd.sh Wan-AI/Wan2.1-T2V-1.3B --tool aria2c -x 10 --hf_token xxxxxxx --hf_username xxxxxxx

3. Clone Repository

git clone git@github.com:Wan-Video/Wan2.1.git

4. Install Dependencies

Note: CUDA >= 11.7 torch >= 2.4.0

cd Wan2.1
pip install -r requirements.txt

5. Generate Videos

5.1. Using the Generate Script

Note: It’s essential to configure –offload_model True –t5_cpu otherwise GPU memory will OOM when saving the video

python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --offload_model True --t5_cpu --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
[2025-02-28 08:33:08,062] INFO: Generation job args: Namespace(task='t2v-1.3B', size='832*480', frame_num=81, ckpt_dir='./Wan2.1-T2V-1.3B', offload_model=True, ulysses_size=1, ring_size=1, t5_fsdp=False, t5_cpu=True, dit_fsdp=False, save_file=None, prompt='Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.', use_prompt_extend=False, prompt_extend_method='local_qwen', prompt_extend_model=None, prompt_extend_target_lang='ch', base_seed=6930324173022001627, image=None, sample_solver='unipc', sample_steps=50, sample_shift=8.0, sample_guide_scale=6.0)
[2025-02-28 08:33:08,063] INFO: Generation model config: {'__name__': 'Config: Wan T2V 1.3B', 't5_model': 'umt5_xxl', 't5_dtype': torch.bfloat16, 'text_len': 512, 'param_dtype': torch.bfloat16, 'num_train_timesteps': 1000, 'sample_fps': 16, 'sample_neg_prompt': '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走', 't5_checkpoint': 'models_t5_umt5-xxl-enc-bf16.pth', 't5_tokenizer': 'google/umt5-xxl', 'vae_checkpoint': 'Wan2.1_VAE.pth', 'vae_stride': (4, 8, 8), 'patch_size': (1, 2, 2), 'dim': 1536, 'ffn_dim': 8960, 'freq_dim': 256, 'num_heads': 12, 'num_layers': 30, 'window_size': (-1, -1), 'qk_norm': True, 'cross_attn_norm': True, 'eps': 1e-06}
[2025-02-28 08:33:08,063] INFO: Input prompt: Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.
[2025-02-28 08:33:08,065] INFO: Creating WanT2V pipeline.
[2025-02-28 08:35:19,682] INFO: loading ./Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth
[2025-02-28 08:37:19,359] INFO: loading ./Wan2.1-T2V-1.3B/Wan2.1_VAE.pth
[2025-02-28 08:37:24,292] INFO: Creating WanModel from ./Wan2.1-T2V-1.3B
[2025-02-28 08:38:17,205] INFO: Generating video ...
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [08:15<00:00,  9.91s/it]
[2025-02-28 08:47:29,986] INFO: Saving generated video to t2v-1.3B_832*480_1_1_Two_anthropomorphic_cats_in_comfy_boxing_gear_and__20250228_084729.mp4
[2025-02-28 08:47:32,176] INFO: Finished.

The video generation is complete. You can find the generated video in the current directory:

t2v-1.3B_832*480_1_1_Two_anthropomorphic_cats_in_comfy_boxing_gear_and__20250228_084729.mp4

Your browser does not support the video tag.

5.2. Using Gradio UI Interface

5.2.1. Launch Gradio Service

This requires a prompt extension API. You can use either the dashscope API or your own API. Here we’ll use the dashscope API https://dashscope.console.aliyun.com/overview

DASH_API_KEY=xxxx  python gradio/t2v_1.3B_singleGPU.py --prompt_extend_method 'dashscope' --ckpt_dir ./Wan2.1-T2V-1.3B/
Step1: Init prompt_expander...done
Step2: Init 1.3B t2v model...done
* Running on local URL:  http://0.0.0.0:7860

To create a public link, set `share=True` in `launch()`.
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [08:14<00:00,  9.88s/it]

5.2.2. Access UI, Input Prompt, and Generate Video

Example prompt:

Over the snow-covered Meili Snow Mountains, multiple fighter jets fly in formation, performing various aerial maneuvers and stunts against a backdrop of clear blue skies and majestic peaks. The scene is dynamic and awe-inspiring.

Your browser does not support the video tag.

HAMI-core 调试

Wed, 22 Jan 2025 00:00:00 +0000

HAMI-core 调试

HAMI-core 调试

HAMI-core 是一个用于拦截CUDA-Runtime(libcudart.so) 发给CUDA-Driver(libcuda.so)请求的动态链接库, 动态库调试相对比较麻烦，这里记录下调试方式

1. 动态链接库编译

1.1. 编译时需要增加debug符号

编译时需要增加debug符号，在CMakeLists.txt中增加-DCMAKE_BUILD_TYPE=Debug 主要标记为 -g -fvisibility=hidden

    set(LIBRARY_COMPILE_FLAGS -shared -fPIC -D_GNU_SOURCE -fvisibility=hidden -Wall)

-g 是生成调试信息
-fvisibility=hidden 是隐藏符号，如果不设置暴露所有方法，在调试时容易碰到glibc版本问题，导致无法调试

/libvgpu/build/test/test_alloc: symbol lookup error: /libvgpu/build/libvgpu.so: undefined symbol: _dl_sym
[Inferior 1 (process 2166963) exited with code 0177]

-Wall 是开启所有警告。

1.2. 设置方法可调试

需要注意 -fvisibility=hidden 会隐藏所有符号，包括 __attribute__((visibility("default"))) 标记的符号，所以需要手动添加 __attribute__((visibility("default"))) 标记的符号。在在项目中 src/libvgpu.c 已定义宏

#define FUNC_ATTR_VISIBLE  __attribute__((visibility("default")))

所以如果需要调试 libvgpu.so 需要手动添加 FUNC_ATTR_VISIBLE 标记即可，例如

FUNC_ATTR_VISIBLE void* dlsym(void* handle, const char* symbol) {
    pthread_once(&dlsym_init_flag,init_dlsym);
    LOG_DEBUG("into dlsym %s",symbol);
    if (real_dlsym == NULL) {
        real_dlsym = dlvsym(RTLD_NEXT,"dlsym","GLIBC_2.2.5");

但是其他文件中未定义 FUNC_ATTR_VISIBLE 宏，所以需要手动添加，需要调试的函数都需要添加，否则调试不到

1.3. 编译

项目中提供了两种方式

make build-in-docker 在docker中构建
make build 在本地构建

有docker的情况下推荐本地在docker中构建，但是需要注意设置项目路径为/libvgpu，否则调试会报错

构建完毕之后在/libvgpu/build目录下会生成如下内容(下面仅列举了后面可能会用到的)

build
├── libvgpu.so
└── test
    ├── test_alloc
    ├── test_alloc_hold
    ├── test_alloc_host
    ├── test_alloc_managed
    ├── test_alloc_pitch
    ├── test_create_3d_array
    ├── test_create_array
    ├── test_host_alloc
    ├── test_host_register
    ├── test_runtime_alloc
    ├── test_runtime_alloc_host
    ├── test_runtime_alloc_managed
    ├── test_runtime_host_alloc
    ├── test_runtime_host_register
    └── test_runtime_launch

2. debug确认

编译完成之后首先要确认 libvgpu.so 是否包含调试信息

readelf -S libvgpu.so | grep debug
# 有如下输出说明有debug信息
  [27] .debug_aranges    PROGBITS         0000000000000000  0004ae2f
  [28] .debug_info       PROGBITS         0000000000000000  0004b0cf
  [29] .debug_abbrev     PROGBITS         0000000000000000  0007016f
  [30] .debug_line       PROGBITS         0000000000000000  00072968
  [31] .debug_str        PROGBITS         0000000000000000  0007c556
  [32] .debug_line_str   PROGBITS         0000000000000000  0008abd6
  [33] .debug_rnglists   PROGBITS         0000000000000000  0008b118

例如我们要调试 nvmlDeviceGetMemoryInfo 方法，且在上面设置了 FUNC_ATTR_VISIBLE 宏，则需要确认 libvgpu.so 中是否包含 nvmlDeviceGetMemoryInfo 方法

nm -D libvgpu.so | grep nvmlDeviceGetMemoryInfo
# 有如下输出即可
0000000000017549 T nvmlDeviceGetMemoryInfo

so侧的设置已经完成，还需要设置gdb侧的设置，例如我们需要以测试用例程序进行gdb则需要检查测试用例是否带了debug标记

readelf -S test/test_alloc | grep debug
# 有如下输出即可
  [28] .debug_aranges    PROGBITS         0000000000000000  0000303b
  [29] .debug_info       PROGBITS         0000000000000000  0000306b
  [30] .debug_abbrev     PROGBITS         0000000000000000  00004ab7
  [31] .debug_line       PROGBITS         0000000000000000  00004db2
  [32] .debug_str        PROGBITS         0000000000000000  00005653
  [33] .debug_line_str   PROGBITS         0000000000000000  0000734c
  [34] .debug_loclists   PROGBITS         0000000000000000  000074a5
  [35] .debug_rnglists   PROGBITS         0000000000000000  000079c6

准备完毕开始设置环境进行调试

3. 动态链接库加载

由于HAMI-core是通过实现 dlsym 来实现动态链接库的加载，所以需要设置环境变量 LD_PRELOAD 指向 libvgpu.so 来实现提前加载，拦截到CUDA-Driver的调用

export LD_PRELOAD=/libvgpu/build/libvgpu.so

验证下

LD_DEBUG=libs ./test/test_alloc
# 输出如下 有calling init: /libvgpu/build/libvgpu.so 说明加载成功
   1758842:	  trying file=/lib/x86_64-linux-gnu/librt.so.1
   1758842:
   1758842:
   1758842:	calling init: /lib64/ld-linux-x86-64.so.2
   1758842:
   1758842:
   1758842:	calling init: /lib/x86_64-linux-gnu/libc.so.6
   1758842:
   1758842:
   1758842:	calling init: /lib/x86_64-linux-gnu/librt.so.1
   1758842:
   1758842:
   1758842:	calling init: /lib/x86_64-linux-gnu/libdl.so.2
   1758842:
   1758842:
   1758842:	calling init: /lib/x86_64-linux-gnu/libm.so.6
   1758842:
   1758842:
   1758842:	calling init: /lib/x86_64-linux-gnu/libpthread.so.0
   1758842:
   1758842:
   1758842:	calling init: /usr/local/cuda-12.4/lib64/libcuda.so.1
   1758842:
   1758842:
   1758842:	calling init: /lib/x86_64-linux-gnu/libnvidia-ml.so.1
   1758842:
   1758842:
   1758842:	calling init: /libvgpu/build/libvgpu.so
   1758842:
   1758842:
   1758842:	initialize program: ./test/test_alloc
   1758842:
   1758842:
   1758842:	transferring control: ./test/test_alloc

4. 调试

gdb ./test/test_alloc
# 打断点, 此处需要看程序内部，例如有main函数，则打断点在main函数
b main
# 运行
r
# 查看so是否被加载
info sharedlibrary
# 有如下输出说明so被加载
# 0x00007ffff7f727f0  0x00007ffff7fa1db0  Yes         /libvgpu/build/libvgpu.so
From                To                  Syms Read   Shared Object Library
0x00007ffff7fc5090  0x00007ffff7fee315  Yes         /lib64/ld-linux-x86-64.so.2
0x00007ffff7f727f0  0x00007ffff7fa1db0  Yes         /libvgpu/build/libvgpu.so
0x00007ffff6c18750  0x00007ffff6d6aae2  Yes (*)     /lib/x86_64-linux-gnu/libnvidia-ml.so.1
0x00007ffff5160d80  0x00007ffff55f75d4  Yes (*)     /usr/local/cuda-12.4/lib64/libcuda.so.1
0x00007ffff4e84700  0x00007ffff501693d  Yes         /lib/x86_64-linux-gnu/libc.so.6
0x00007ffff7f5f040  0x00007ffff7f5f105  Yes         /lib/x86_64-linux-gnu/libpthread.so.0
0x00007ffff7e833a0  0x00007ffff7efe8c8  Yes         /lib/x86_64-linux-gnu/libm.so.6
0x00007ffff7e71040  0x00007ffff7e71105  Yes         /lib/x86_64-linux-gnu/libdl.so.2
0x00007ffff7e6c080  0x00007ffff7e6c275  Yes         /lib/x86_64-linux-gnu/librt.so.1
# 在so的需要调试的函数打断点
b nvmlDeviceGetMemoryInfo
# 确认断点位置正确
list nvmlDeviceGetMemoryInfo
# 输出如下, 符合在代码中的预期
362	    }
363	    return NVML_SUCCESS;
364	}
365
366	 __attribute__((visibility("default")))
367	 nvmlReturn_t nvmlDeviceGetMemoryInfo(nvmlDevice_t device, nvmlMemory_t* memory) {
368	    return _nvmlDeviceGetMemoryInfo(device,memory,1);
369	}
370
371	nvmlReturn_t nvmlDeviceGetMemoryInfo_v2(nvmlDevice_t device, nvmlMemory_v2_t* memory) {
# 或者
(gdb) info functions  nvmlDeviceGetMemoryInfo
# 输出如下 说明断点在 `/libvgpu/src/nvml/hook.c` 文件中, 拦截成功符合预期
All functions matching regular expression "nvmlDeviceGetMemoryInfo":

File /libvgpu/src/nvml/hook.c:
317:	nvmlReturn_t _nvmlDeviceGetMemoryInfo(nvmlDevice_t, nvmlMemory_t *, int);
367:	nvmlReturn_t nvmlDeviceGetMemoryInfo(nvmlDevice_t, nvmlMemory_t *);
371:	nvmlReturn_t nvmlDeviceGetMemoryInfo_v2(nvmlDevice_t, nvmlMemory_v2_t *);

Non-debugging symbols:
0x0000555555555180  nvmlDeviceGetMemoryInfo@plt
0x00007ffff6c18580  nvmlDeviceGetMemoryInfo_v2@plt
0x00007ffff6c494a0  nvmlDeviceGetMemoryInfo
0x00007ffff6c496d0  nvmlDeviceGetMemoryInfo_v2
# 运行到断点
c
# 查看断点栈信息
bt
# 查看所有断点
info b
# 输出如下，符合预期
0  nvmlDeviceGetMemoryInfo (device=0x7ffff6ffaf38, memory=0x7fffffffdfc0) at /libvgpu/src/nvml/hook.c:368
1  0x00005555555553d8 in get_current_memory_usage (usage=usage@entry=0x555555558030 <usage>) at /libvgpu/test/test_utils.h:70
2  0x0000555555555831 in main () at /libvgpu/test/test_alloc.c:76
# 如果错过断点，可以重新来
r
# 列出所有可以debug的方法名
info functions
# 输出如下
All defined functions:

Non-debugging symbols:
0x0000000000402f10  getenv@plt
0x0000000000402f20  free@plt
0x0000000000402f30  putchar@plt
0x0000000000402f40  pthread_create@plt
0x0000000000402f50  endmntent@plt
0x0000000000402f60  strcasecmp@plt
# 设置变量
set {char[5]} 0x7fffffffdfc0 = "nvmlInit"
# 查看变量
print 0x7fffffffdfc0
# 输出如下
$1 = "nvmlInit"
# 设置环境变量
set environment LD_PRELOAD=/libvgpu/build/libvgpu.so
# 查看环境变量
show environment LD_PRELOAD

5. 一个debug例子

HAMI-core 项目中 test/test_alloc.c 文件中 main 函数中调用了 nvmlDeviceGetMemoryInfo 方法, 我们debug下这个方法

通过读代码可知，调用流程为首先调用dlsym, dlsym会初始化方法，这里nvml使用 nvmlInitWithFlags 注册nvml重写方法，调用时会找到注册的方法地址，然后通过方法地址调用重写的方法，所以debug时首先需要找到dlsym的调用位置，然后找到重写的方法地址，最后找到重写的方法调用位置

export LD_PRELOAD=/libvgpu/build/libvgpu.so
gdb ./test/test_alloc
# 打断点
b main
# 运行
r
# 打dlsym的断点，这里会有多个实现所以可以指定文件打断点
# 运行到这里的时候需要手动更改下symbol的值和环境变量HOOK_NVML_ENABLE的值，确保重写的方法nvmlDeviceGetMemoryInfo内存地址被hook到
b /libvgpu/src/libvgpu.c:77
# 找到nvmlDeviceGetMemoryInfo方法
info functions nvmlDeviceGetMemoryInfo
# 打断点
b /libvgpu/src/nvml/hook.c:368
# 在进入下一层
b get_current_device_memory_usage
# 开始运行
c
# 输出如下
(gdb) c
Continuing.
[New Thread 0x7ffff4e5b640 (LWP 2200336)]

Thread 1 "test_alloc" hit Breakpoint 2, dlsym (handle=0x7ffff7ffe2e0, symbol=0x7ffff69f26c0 "dladdr1") at /libvgpu/src/libvgpu.c:78
78	    pthread_once(&dlsym_init_flag,init_dlsym);
# 这里需要修改 symbol为 nvmlInitWithFlags
set {char[18]} 0x7ffff69f26c0 = "nvmlInitWithFlags"
# 验证
p symbol
$2 = 0x7ffff69f26c0 "nvmlDeviceGetMemoryInfo"
# 设置环境变量 HOOK_NVML_ENABLE 为 true
set environment HOOK_NVML_ENABLE = true
# 验证
show environment HOOK_NVML_ENABLE
HOOK_NVML_ENABLE = true
# 下一个断点
(gdb) c
Continuing.
[New Thread 0x7fffe8fde640 (LWP 2219816)]
[New Thread 0x7fffe2fde640 (LWP 2219817)]

Thread 1 "test_alloc" hit Breakpoint 5, nvmlDeviceGetMemoryInfo (device=0x7ffff6ffaf38, memory=0x7fffffffdfc0) at /libvgpu/src/nvml/hook.c:367
367	    return _nvmlDeviceGetMemoryInfo(device,memory,1);
(gdb) s
_nvmlDeviceGetMemoryInfo (device=0x7ffff6ffaf38, memory=0x7fffffffdfc0, version=1) at /libvgpu/src/nvml/hook.c:317
317	FUNC_ATTR_VISIBLE nvmlReturn_t _nvmlDeviceGetMemoryInfo(nvmlDevice_t device,nvmlMemory_t* memory,int version) {
(gdb) l
312	}
313
314	void nvml_postInit() {
315	}
316
317	FUNC_ATTR_VISIBLE nvmlReturn_t _nvmlDeviceGetMemoryInfo(nvmlDevice_t device,nvmlMemory_t* memory,int version) {
318	    unsigned int dev_id;
319	    LOG_DEBUG("into nvmlDeviceGetMemoryInfo");
320
321	    switch (version){
(gdb) c
Continuing.

Thread 1 "test_alloc" hit Breakpoint 4, get_current_device_memory_usage (dev=0) at /libvgpu/src/multiprocess/multiprocess_memory_limit.c:877

这样我们就可以debug整个链路

6. 例子2: debug nvidia-smi

export LD_PRELOAD=/libvgpu/build/libvgpu.so
whereis nvidia-smi
gdb /usr/bin/nvidia-smi
# 找到可以打断点的方法
info functions
# 例如dlysm是一定会被调用的
b dlsym
# 开始运行
r
# run之后动态库被加载，可以debug动态库中的方法
# 查看哪些方法可以debug
info functions
# 例如nvidia-smi一定会调用到get_gpu_memory_usage
b get_gpu_memory_usage
# 跳到断点
c
# 这里会反复循环
Breakpoint 1, ___dlsym (handle=0x718f30, name=0x7ffff7fa8520 "nvmlVgpuInstanceSetEncoderCapacity") at ./dlfcn/dlsym.c:64
64	in ./dlfcn/dlsym.c
# 是因为上面的dlsym有多个实现，这里我们不需要知道系统的dlsym是怎么实现的所以可以把这个断点删掉
info b
# 输出如下
Num     Type           Disp Enb Address            What
2       breakpoint     keep y   0x00007ffff4e695b4 <nvmlDeviceGetNumaNodeId+4>
4       breakpoint     keep y   <MULTIPLE>
4.1                         y   0x00007ffff4e55ec4 <nvmlVgpuInstanceGetVmID+4>
4.2                         y   0x00007ffff7f81027 in nvmlVgpuInstanceGetVmID at /libvgpu/src/nvml/nvml_entry.c:1132
5       breakpoint     keep y   0x00007ffff7f9cf5a in get_gpu_memory_usage at /libvgpu/src/multiprocess/multiprocess_memory_limit.c:291
6       breakpoint     keep y   <MULTIPLE>
	breakpoint already hit 4 times
6.1                         y   0x00007ffff7cd9740 in ___dlsym at ./dlfcn/dlsym.c:64
6.2                         y   0x00007ffff7f72acf in dlsym at /libvgpu/src/libvgpu.c:78
# 禁用6.1这个断点
disable 6.1
# 继续就可以debug到了
c
Breakpoint 6, dlsym (handle=0x718f30, symbol=0x4d1e79 "nvmlDeviceGetMemoryInfo_v2") at /libvgpu/src/libvgpu.c:78
78	    pthread_once(&dlsym_init_flag,init_dlsym);
(gdb)
Continuing.

Breakpoint 5, get_gpu_memory_usage (dev=0) at /libvgpu/src/multiprocess/multiprocess_memory_limit.c:291
291	    LOG_INFO("get_gpu_memory_usage dev=%d",dev);
(gdb)

7. 使用clion远程debug

7.1. 代码同步

Tools -> Deployment -> Configuration

添加一个SFTP, 填写玩connection之后填写目录映射

Local path: /go/src/github.com/Project-HAMi/HAMi-core

Remote path: /libvgpu

勾选同步 Use Rsync for download/upload/sync

同步代码

有两种方式，一种是gdb远程调试，另一种是clion远程调试

7.2. gdb远程调试

Run/Debug -> Debugger -> GDB Remote Debugging -> 填写Remote Host和Remote Port

'target remote' args: tcp:10.10.200.8:1234

Path Mapping: /libvgpu -> /go/src/github.com/Project-HAMi/HAMi-core

在服务器上启动gdbserver

gdbserver :1234 ./test/test_alloc

本地打断点，点击debug按钮

首先断点dlsym的调用位置,修改symbol为nvmlInitWithFlags

设置环境变量HOOK_NVML_ENABLE为true

继续断点，一直到最终的 get_gpu_memory_usage

7.3. clion远程调试

Clion -> Settings -> Build, Execution, Deployment -> Toolchains 新增一个

Clion -> Settings -> Build, Execution, Deployment -> CMake 新增一个

build type选择 debug，toolchain选择刚刚创建的，其他默认即可

配置完毕之后即可在 Run/Debug Configuration 看到许多可以调试的应用

选择一个调试即可，例如调试test_alloc则需配置 executable 和 Environment variable这样可直接调试，无需手动在服务器上启动gdbserver

Golang Interview

Tue, 14 Jan 2025 00:00:00 +0000

知识点

goroutine为什么比线程快

创建调度与销毁: goroutine的创建调度与销毁都在用户态，线程的在内核态
内存分配: goroutine栈内存为2-4KB，线程1MB
切换: goroutine切换时仅需要保护少数寄存器，而线程则需要保护大量的寄存器，goroutine无需操作系统线程切换

GMP

两级线程模型，内核调度实体kse与线程的对应关系为m:n
kse -> m -> p ->g
- m: machine 关联一个内核线程
- p: processor 作为上下文，衔接m和g
- g: goroutine 轻量化线程，调用栈、调用信息，例如channel等

https://github.com/fengyuan-liang/notes/blob/main/GoLang/golang%E5%A4%A7%E6%9D%80%E5%99%A8GMP%E6%A8%A1%E5%9E%8B.md

https://www.jianshu.com/p/4afa0679851d

slice的扩容

golang 1.21.4

对于小slice进行1.25倍扩展逐渐过度
对于大slice进行2倍扩展

src/runtime/slice.go

 newcap := oldCap
 doublecap := newcap + newcap
 if newLen > doublecap {
  newcap = newLen
 } else {
  const threshold = 256
  if oldCap < threshold {
   newcap = doublecap
  } else {
   // Check 0 < newcap to detect overflow
   // and prevent an infinite loop.
   for 0 < newcap && newcap < newLen {
    // Transition from growing 2x for small slices
    // to growing 1.25x for large slices. This formula
    // gives a smooth-ish transition between the two.
    newcap += (newcap + 3*threshold) / 4
   }
   // Set newcap to the requested cap when
   // the newcap calculation overflowed.
   if newcap <= 0 {
    newcap = newLen
   }
  }
 }

g0

持有调度栈的goroutine，作用有

goroutine的创建
大内存分配
CGO函数的执行

三色标记法

https://community.apinto.com/d/34057-golang-gc

Data Structures and

Sliding Window(滑动窗口)

第三题

https://leetcode.cn/problems/longest-substring-without-repeating-characters/

关键点:

index>=left: index一定需要在left右边，保证取值不会超过当前窗口
left=index+1: 如果前面的有重复了说明当前窗口已经是最大状态，从index+1开始才能保证和当前窗口无重复字母
index的更新: 有重复字母则更新对应的坐标

func lengthOfLongestSubstring(s string) int {
  charMap := make(map[byte]int)
  maxLength := 0
  left, right := 0,0
  for right=0;right < len(s);right++ {
    if index, ok := charMap[s[right]]; ok && index >= left {
      left = index + 1
    }
    charMap[s[right]] = right
    maxLength = max(maxLength, right-left+1)
  }
  return maxLength
}

Brute Force(暴力枚举)

递归（Recursion）

回溯（Backtracking）

分治（Divide and Conquer）

动态规划（Dynamic Programming, DP）

第五题，最大回文子串

https://leetcode.cn/problems/longest-palindromic-substring/description/

关键点:

先算小的，小的算完之后记录结果, 用小的结果计算更大的记录结果，然后得到大的结果
- 单个字符一定是回文 dp[i][i]=true
- 相邻字符如果是相同字符一定是回文 dp[i][i+1]=true
- 大于三个的则可以通过 dp[i][j] 如果头尾一样 s[i] == s[j], 那么判断 dp[i+1][j-1] 为true则为true
- 那么通过从小到多的遍历就能获取到所欲呕的dp[i][j]
- 第一次遍历, 处理只有一位的回文
- 第二次遍历，填充二位回文
- 第三次遍历，填充3位及以上回文，这里需要注意
  - 由于前两次遍历已处理完长度为1和长度为2的回文字符串，所以遍历时字符串长度从3开始，进行遍历所有大于3的长度的字符串填充
  - 有了字符串长度的一重循环，内部需要从第一个字符开始遍历，最大遍历到字符串总体长度-当前计算的字符串长度，注意边界问题
    - 例如一个字符串长度为9，当前计算的长度为5 那么最后一个坐标应该是 9-5=4，可以取到4吗 0 1 2 3 4 5 6 7 8，取到4 0 1 2 3 + 后面五位 4 5 6 7 8 9 所以这里可以取到这个边界值
    - 二重循环是有了起点还需要有一个终点，这个终点就是起点坐标+长度-1 保证不会超出整个字符串的长度

func longestPalindrome(s string) string {
 n := len(s)

 if n == 0 {
  return ""
 }

 dp := make([][]bool, n)
 for i := range dp {
  dp[i] = make([]bool, n)
 }
 start, maxLen := 0, 1
 for i := 0; i < n; i++ {
  dp[i][i] = true
 }
 for i := 0; i < n-1; i++ {
  if s[i] == s[i+1] {
   dp[i][i+1] = true
   start = i
   maxLen = 2
  }
 }

 for length := 3; length <= n; length++ {
  for i := 0; i <= n-length; i++ {
   j := i + length - 1
   if s[i] == s[j] && dp[i+1][j-1] {
    dp[i][j] = true
    start = i
    maxLen = length
   }
  }
 }
 return s[start : start+maxLen]

贪心算法（Greedy Algorithm）

双指针（Two Pointers）

并查集（Union-Find）

拓扑排序（Topological Sorting）

图算法（Graph Algorithms）

位运算（Bit Manipulation）

单调栈 & 单调队列（Monotonic Stack & Queue）

Algorithms

LRU

第146题

https://leetcode.cn/problems/lru-cache/submissions/609910341/

package lrucache

type LRUCache struct {
 cache            map[int]*Node
 capacity, length int
 head, tail       *Node
}

type Node struct {
 key, value int
 next, prev *Node
}

func Constructor(capacity int) LRUCache {
 head, tail := &Node{}, &Node{}
 head.next = tail
 tail.prev = head
 return LRUCache{
  cache:    make(map[int]*Node, capacity),
  capacity: capacity,
  head:     head,
  tail:     tail,
 }
}

func (this *LRUCache) Get(key int) int {
 if node, ok := this.cache[key]; ok {
  this.movetoHead(node)
  return node.value
 }
 return -1

}

func (this *LRUCache) Put(key int, value int) {
 node, ok := this.cache[key]
 if ok {
  node.value = value
  this.movetoHead(node)
 } else {
  if this.length >= this.capacity {
   this.removeTail()
  }
  node = &Node{key: key, value: value}
  this.addToHead(node)
  this.cache[key] = node
  this.length++
 }
}

func (this *LRUCache) addToHead(node *Node) {
 this.head.next.prev = node
 node.next = this.head.next
 this.head.next = node
 node.prev = this.head

}

func (this *LRUCache) removeTail() {
 if this.tail.prev == this.head {
  return
 }
 node := this.tail.prev
 this.removeNode(node)
 delete(this.cache, node.key)
 this.length--
}

func (this *LRUCache) movetoHead(node *Node) {
 this.removeNode(node)
 this.addToHead(node)
}

func (this *LRUCache) removeNode(node *Node) {
 node.prev.next = node.next
 node.next.prev = node.prev
}

/**
 * Your LRUCache object will be instantiated and called as such:
 * obj := Constructor(capacity);
 * param_1 := obj.Get(key);
 * obj.Put(key,value);
 */

HAMI项目GPU POD调度流程源码走读

Tue, 24 Dec 2024 00:00:00 +0000

1. 概述
2. 调度流程
3. Pod调度流程
4. 参考

1. 概述

使用hami的过程中经常会出现Pod被创建出来Pending的问题，犹以如下两个问题为著

Pod UnexpectedAdmissionError
Pod Pending

介于此，展开这部分代码的粗略走读，旨在说明调度过程中各组件的交互，以及资源的计算方式，其他细节会有所遗漏

2. 调度流程

看代码之前可以先看下官方文档说明，大体上比较明确

细节上可以分为三个阶段

准备阶段: 图上可以看出有一些依赖条件，例如要有mutating webhook， device-plugin等等，所以这个阶段主要分析下依赖条件的准备，只有在服务首次启动时需要

pod调度阶段: 准备过程完成之后pod进入处理流程，完成调度
Pod启动阶段: pod如何与node上的gpu进行交互等

本文会着重分析准备阶段，主要内容为调度分析

3. Pod调度流程

用户发送创建Pod请求到kube-apiserver
触发adminssion webhook，更新pod中schedulerName
kube-apiserver根据schedulerName将请求发送给调度器处理
调度器处理
- 收集Node device信息 – 通过node annotation收集，数据来自daemonSet hami-device-plugin 定时写入
- 根据设备信息以及pod的limit信息进行打分，选出最高分的node
- 将pod和node进行绑定完成绑定，进行pod创建

常见的几个问题排查及处理

Pod UnexpectedAdmissionError

pod创建状态显示 UnexpectedAdmissionError

了解流程之后，可以知道这个错误代表kube-apiserver调用拓展调度器失败，可能有两个原因，其他情况具体排查需要看 kube-apiserver日志

通信异常：从kube-apiserver到拓展调度器的https端口不通，有几种可能
- dns无法解析
- 跨节点通信有问题
- 拓展调度器的服务异常
TLS验证错误: 一般会显示 webhook x509: certificate signed by unknown authority，helmchart部署时有一个 jobs.batch hami-vgpu.admission-pathch, 如果没有运行完成会出现这样的问题

调度问题

容器一直在 pending 状态，使用 kubectl describe命令可以看到具体原因，主要有以下几个

card Insufficient remaining memory
calcScore:node not fit pod

主要原因一般会是确实资源不足，或者配置错误，配置错误是指 devicememoryscaling 配置未符合预期，有两个地方可以配置, 优先级为节点配置大于全局配置，容易发生问题的地方在于 name需要和kubectl get node显示的nodename一致才能生效

全局配置 k get cm hami-scheduler-device

  deviceMemoryScaling: 3

节点配置 k get cm hami-device-plugin

{
  "nodeconfig": [
    {
      "name": "node1",
      "devicememoryscaling": 3,
      "devicesplitcount": 10,
      "migstrategy":"none",
      "filterdevices": {
        "uuid": [],
        "index": []
      }
    }
  ]
}

3.1. MutatingWebhook

K8s提供了adminssionWebhook资源, 以k8s资源操作为触发器，触发hook，用途最广泛的为针对Pod创建做拦截，对Pod做yaml注入，具体的例如增加init容器注入文件等等

3.1.1. webhook配置

hami-webhook

k get mutatingwebhookconfigurations.admissionregistration.k8s.io hami-webhook -o yaml
apiVersion: admissionregistration.k8s.io/v1
kind: MutatingWebhookConfiguration
metadata:
  annotations:
    meta.helm.sh/release-name: hami
    meta.helm.sh/release-namespace: kube-system
  creationTimestamp: "2024-12-10T03:50:37Z"
  generation: 5
  labels:
    app.kubernetes.io/managed-by: Helm
  name: hami-webhook
  resourceVersion: "2307810"
  uid: 2cdcebe4-f561-429f-9480-701e65980687
webhooks:
- admissionReviewVersions:
  - v1beta1
  clientConfig:
    caBundle: LS0tLS1CRUdJTiBDRVJUSUZJQ0FURS0tLS0tCk1JSUJkakNDQVJ5Z0F3SUJBZ0lSQUxjd2FQMjUrMlphdGhTTlFMcG1qT0V3Q2dZSUtvWkl6ajBFQXdJd0R6RU4KTUFzR0ExVUVDaE1FYm1sc01UQWdGdzB5TkRFeU1EWXdOekV4TVRWYUdBOHlNVEkwTVRFeE1qQTNNVEV4TlZvdwpEekVOTUFzR0ExVUVDaE1FYm1sc01UQlpNQk1HQnlxR1NNNDlBZ0VHQ0NxR1NNNDlBd0VIQTBJQUJDUnlXUDdYCkRmT2N4NEVTMVRYaUs0dnFFU2wrcUFHYjI2YzNrOEdMWlZTL1lHaFpLZVVxaEgydVRhTFdWTW1hZVJFbkxqM0cKSStMVFRVTTR6SVhEUld5alZ6QlZNQTRHQTFVZER3RUIvd1FFQXdJQ0JEQVRCZ05WSFNVRUREQUtCZ2dyQmdFRgpCUWNEQVRBUEJnTlZIUk1CQWY4RUJUQURBUUgvTUIwR0ExVWREZ1FXQkJTcVV4bWpGa29YUlpRK0xXVzBNM1pJCnMzck1wakFLQmdncWhrak9QUVFEQWdOSUFEQkZBaUJSY2VRL2tJVkR2VTV3Vjl0K3NRWm93TmFhTWhIMTV5K2sKT3VrR0FlRGVtQUloQUxDZzFrM0JQZUJBNG8reWY5emxvVjM2VEk2RHUzaGdMT1B3MXhaZkFvcDMKLS0tLS1FTkQgQ0VSVElGSUNBVEUtLS0tLQo=
    service:
      name: hami-scheduler
      namespace: kube-system
      path: /webhook
      port: 443
  failurePolicy: Ignore
  matchPolicy: Equivalent
  name: vgpu.hami.io
  namespaceSelector:
    matchExpressions:
    - key: hami.io/webhook
      operator: NotIn
      values:
      - ignore
  objectSelector:
    matchExpressions:
    - key: hami.io/webhook
      operator: NotIn
      values:
      - ignore
  reinvocationPolicy: Never
  rules:
  - apiGroups:
    - ""
    apiVersions:
    - v1
    operations:
    - CREATE
    resources:
    - pods
    scope: '*'
  sideEffects: None
  timeoutSeconds: 10

当pod创建时，调用 https://hami-scheduler.kube-system:443/webhook 做tls校验，CA为 caBundle 配置当命名空间有 hami.io/webhook: ignore 的标签时不触发

3.1.2. webhook Server实现

需要实现一个tls的Http Server, 且提供 /webhook 接口

cmd/scheduler/main.go:84

func start() {
 ...
 router.POST("/webhook", routes.WebHookRoute())

WebHookRoute 需要实现 sigs.k8s.io/controller-runtime@v0.16.3/pkg/webhook/admission/webhook.go:98

pkg/scheduler/webhook.go:52

 pod := &corev1.Pod{}
 err := h.decoder.Decode(req, pod)
 if err != nil {
  klog.Errorf("Failed to decode request: %v", err)
  return admission.Errored(http.StatusBadRequest, err)
 }
 if len(pod.Spec.Containers) == 0 {
  klog.Warningf(template+" - Denying admission as pod has no containers", req.Namespace, req.Name, req.UID)
  return admission.Denied("pod has no containers")
 }
 klog.Infof(template, req.Namespace, req.Name, req.UID)
 hasResource := false
 for idx, ctr := range pod.Spec.Containers {
  c := &pod.Spec.Containers[idx]
  if ctr.SecurityContext != nil {
   if ctr.SecurityContext.Privileged != nil && *ctr.SecurityContext.Privileged {
    klog.Warningf(template+" - Denying admission as container %s is privileged", req.Namespace, req.Name, req.UID, c.Name)
    continue
   }
  }
  for _, val := range device.GetDevices() {
   found, err := val.MutateAdmission(c, pod)
   if err != nil {
    klog.Errorf("validating pod failed:%s", err.Error())
    return admission.Errored(http.StatusInternalServerError, err)
   }
   hasResource = hasResource || found
  }
 }

 if !hasResource {
  klog.Infof(template+" - Allowing admission for pod: no resource found", req.Namespace, req.Name, req.UID)
  //return admission.Allowed("no resource found")
 } else if len(config.SchedulerName) > 0 {
  pod.Spec.SchedulerName = config.SchedulerName
  if pod.Spec.NodeName != "" {
   klog.Infof(template+" - Pod already has node assigned", req.Namespace, req.Name, req.UID)
   return admission.Denied("pod has node assigned")
  }
 }
 marshaledPod, err := json.Marshal(pod)
 if err != nil {
  klog.Errorf(template+" - Failed to marshal pod, error: %v", req.Namespace, req.Name, req.UID, err)
  return admission.Errored(http.StatusInternalServerError, err)
 }
 return admission.PatchResponseFromRaw(req.Object.Raw, marshaledPod)

主要通过pod中容器的resource来判断是否要不要走拓展调度器

pkg/device/nvidia/device.go:246

func (dev *NvidiaGPUDevices) MutateAdmission(ctr *corev1.Container, p *corev1.Pod) (bool, error) {
 /*gpu related */
 priority, ok := ctr.Resources.Limits[corev1.ResourceName(dev.config.ResourcePriority)]
 if ok {
  ctr.Env = append(ctr.Env, corev1.EnvVar{
   Name:  util.TaskPriority,
   Value: fmt.Sprint(priority.Value()),
  })
 }

 _, resourceNameOK := ctr.Resources.Limits[corev1.ResourceName(dev.config.ResourceCountName)]
 if resourceNameOK {
  return resourceNameOK, nil
 }

 _, resourceCoresOK := ctr.Resources.Limits[corev1.ResourceName(dev.config.ResourceCoreName)]
 _, resourceMemOK := ctr.Resources.Limits[corev1.ResourceName(dev.config.ResourceMemoryName)]
 _, resourceMemPercentageOK := ctr.Resources.Limits[corev1.ResourceName(dev.config.ResourceMemoryPercentageName)]

 if resourceCoresOK || resourceMemOK || resourceMemPercentageOK {
  if dev.config.DefaultGPUNum > 0 {
   ctr.Resources.Limits[corev1.ResourceName(dev.config.ResourceCountName)] = *resource.NewQuantity(int64(dev.config.DefaultGPUNum), resource.BinarySI)
   resourceNameOK = true
  }
 }

 if !resourceNameOK && dev.config.OverwriteEnv {
  ctr.Env = append(ctr.Env, corev1.EnvVar{
   Name:  "NVIDIA_VISIBLE_DEVICES",
   Value: "none",
  })
 }
 return resourceNameOK, nil
}

主要比对pod的 Resources Limit中有没有包含 device-config.yaml 的配置，如果有走hami调度流程

deivce-config 以英伟达显卡为例

nvidia:
  resourceCountName: nvidia.com/gpu
  resourceMemoryName: nvidia.com/gpumem
  resourceMemoryPercentageName: nvidia.com/gpumem-percentage
  resourceCoreName: nvidia.com/gpucores
  resourcePriorityName: nvidia.com/priority
  overwriteEnv: false
  defaultMemory: 0
  defaultCores: 0
  defaultGPUNum: 1
  deviceSplitCount: 10
  deviceMemoryScaling: 3
  deviceCoreScaling: 3

确定走hami调度流程之后，通过Patch修改Pod schedulerName 为hami调度器的名称

3.2. 拓展k8s scheduler

KubeSchedulerConfiguration 拓展调度器可以通过实现拓展点进行调度器的拓展

3.2.1. KubeSchedulerConfiguration

k get cm hami-scheduler-newversion -o yaml
apiVersion: v1
data:
  config.yaml: |
    apiVersion: kubescheduler.config.k8s.io/v1beta2
    kind: KubeSchedulerConfiguration
    leaderElection:
      leaderElect: false
    profiles:
    - schedulerName: hami-scheduler
    extenders:
    - urlPrefix: "https://127.0.0.1:443"
      filterVerb: filter
      bindVerb: bind
      nodeCacheCapable: true
      weight: 1
      httpTimeout: 30s
      enableHTTPS: true
      tlsConfig:
        insecure: true
      managedResources:
      - name: nvidia.com/gpu
        ignoredByScheduler: true
      - name: nvidia.com/gpumem
        ignoredByScheduler: true
      - name: nvidia.com/gpucores
        ignoredByScheduler: true
      - name: nvidia.com/gpumem-percentage
        ignoredByScheduler: true
      - name: nvidia.com/priority
        ignoredByScheduler: true
      - name: cambricon.com/vmlu
        ignoredByScheduler: true
      - name: hygon.com/dcunum
        ignoredByScheduler: true
      - name: hygon.com/dcumem
        ignoredByScheduler: true
      - name: hygon.com/dcucores
        ignoredByScheduler: true
      - name: iluvatar.ai/vgpu
        ignoredByScheduler: true
kind: ConfigMap
metadata:
  annotations:
    meta.helm.sh/release-name: hami
    meta.helm.sh/release-namespace: kube-system
  creationTimestamp: "2024-12-10T03:50:36Z"
  labels:

    app.kubernetes.io/component: hami-scheduler
    app.kubernetes.io/instance: hami
    app.kubernetes.io/managed-by: Helm
    app.kubernetes.io/name: hami
    app.kubernetes.io/version: 2.4.1
    helm.sh/chart: hami-2.4.1
  name: hami-scheduler-newversion
  namespace: kube-system
  resourceVersion: "2316275"
  uid: 3a61a72c-0bab-432f-b4d7-5c1ae46ee14d

拓展调度器通过拓展点进行拓展, 这里拓展了filter和bind

filter: 找到最合适的node
bind: 为pod创建一个binding资源

调度时会根据拓展点顺序来调用拓展调度器的实现，这里会先调用

https://127.0.0.1:443/filter 再调用 https://127.0.0.1:443/filter

3.2.2. 拓展调度器Http Server启动

cmd/scheduler/main.go:70

func start() {
 device.InitDevices()
 sher = scheduler.NewScheduler()
 sher.Start()
 defer sher.Stop()

 // start monitor metrics
 go sher.RegisterFromNodeAnnotations()
 go initMetrics(config.MetricsBindAddress)

 // start http server
 router := httprouter.New()
 router.POST("/filter", routes.PredicateRoute(sher))
 router.POST("/bind", routes.Bind(sher))

3.2.3. filter实现

pkg/scheduler/routes/route.go:41

func PredicateRoute(s *scheduler.Scheduler) httprouter.Handle {
 klog.Infoln("Into Predicate Route outer func")
 return func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
  klog.Infoln("Into Predicate Route inner func")
  checkBody(w, r)

  var buf bytes.Buffer
  body := io.TeeReader(r.Body, &buf)

  var extenderArgs extenderv1.ExtenderArgs
  var extenderFilterResult *extenderv1.ExtenderFilterResult

  if err := json.NewDecoder(body).Decode(&extenderArgs); err != nil {
   klog.Errorln("decode error", err.Error())
   extenderFilterResult = &extenderv1.ExtenderFilterResult{
    Error: err.Error(),
   }
  } else {
   extenderFilterResult, err = s.Filter(extenderArgs)
   if err != nil {
    klog.Errorf("pod %v filter error, %v", extenderArgs.Pod.Name, err)
    extenderFilterResult = &extenderv1.ExtenderFilterResult{
     Error: err.Error(),
    }
   }
  }

  if resultBody, err := json.Marshal(extenderFilterResult); err != nil {
   klog.Errorf("Failed to marshal extenderFilterResult: %+v, %+v",
    err, extenderFilterResult)
   w.Header().Set("Content-Type", "application/json")
   w.WriteHeader(http.StatusInternalServerError)
   w.Write([]byte(err.Error()))
  } else {
   w.Header().Set("Content-Type", "application/json")
   w.WriteHeader(http.StatusOK)
   w.Write(resultBody)
  }
 }
}

pkg/scheduler/scheduler.go:430

func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFilterResult, error) {
 klog.InfoS("begin schedule filter", "pod", args.Pod.Name, "uuid", args.Pod.UID, "namespaces", args.Pod.Namespace)
 nums := k8sutil.Resourcereqs(args.Pod)
 total := 0
 for _, n := range nums {
  for _, k := range n {
   total += int(k.Nums)
  }
 }
 if total == 0 {
  klog.V(1).Infof("pod %v not find resource", args.Pod.Name)
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, fmt.Errorf("does not request any resource"))
  return &extenderv1.ExtenderFilterResult{
   NodeNames:   args.NodeNames,
   FailedNodes: nil,
   Error:       "",
  }, nil
 }
 annos := args.Pod.Annotations
 s.delPod(args.Pod)
 nodeUsage, failedNodes, err := s.getNodesUsage(args.NodeNames, args.Pod)
 if err != nil {
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, err)
  return nil, err
 }
 if len(failedNodes) != 0 {
  klog.V(5).InfoS("getNodesUsage failed nodes", "nodes", failedNodes)
 }
 nodeScores, err := s.calcScore(nodeUsage, nums, annos, args.Pod)
 if err != nil {
  err := fmt.Errorf("calcScore failed %v for pod %v", err, args.Pod.Name)
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, err)
  return nil, err
 }
 if len((*nodeScores).NodeList) == 0 {
  klog.V(4).Infof("All node scores do not meet for pod %v", args.Pod.Name)
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, fmt.Errorf("no available node, all node scores do not meet"))
  return &extenderv1.ExtenderFilterResult{
   FailedNodes: failedNodes,
  }, nil
 }
 klog.V(4).Infoln("nodeScores_len=", len((*nodeScores).NodeList))
 sort.Sort(nodeScores)
 m := (*nodeScores).NodeList[len((*nodeScores).NodeList)-1]
 klog.Infof("schedule %v/%v to %v %v", args.Pod.Namespace, args.Pod.Name, m.NodeID, m.Devices)
 annotations := make(map[string]string)
 annotations[util.AssignedNodeAnnotations] = m.NodeID
 annotations[util.AssignedTimeAnnotations] = strconv.FormatInt(time.Now().Unix(), 10)

 for _, val := range device.GetDevices() {
  val.PatchAnnotations(&annotations, m.Devices)
 }

 //InRequestDevices := util.EncodePodDevices(util.InRequestDevices, m.devices)
 //supportDevices := util.EncodePodDevices(util.SupportDevices, m.devices)
 //maps.Copy(annotations, InRequestDevices)
 //maps.Copy(annotations, supportDevices)
 s.addPod(args.Pod, m.NodeID, m.Devices)
 err = util.PatchPodAnnotations(args.Pod, annotations)
 if err != nil {
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, err)
  s.delPod(args.Pod)
  return nil, err
 }
 s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringSucceed, []string{m.NodeID}, nil)
 res := extenderv1.ExtenderFilterResult{NodeNames: &[]string{m.NodeID}}
 return &res, nil
}

这里核心逻辑主要有两步, 获取节点资源、根据节点已分配资源与总资源计算分数并选出一个最高分

3.2.3.1. 获取节点资源信息

pkg/scheduler/scheduler.go:241

func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[string]*NodeUsage, map[string]string, error) {
 overallnodeMap := make(map[string]*NodeUsage)
 cachenodeMap := make(map[string]*NodeUsage)
 failedNodes := make(map[string]string)
 //for _, nodeID := range *nodes {
 allNodes, err := s.ListNodes()
 if err != nil {
  return &overallnodeMap, failedNodes, err
 }

 for _, node := range allNodes {
  nodeInfo := &NodeUsage{}
  userGPUPolicy := config.GPUSchedulerPolicy
  if task != nil && task.Annotations != nil {
   if value, ok := task.Annotations[policy.GPUSchedulerPolicyAnnotationKey]; ok {
    userGPUPolicy = value
   }
  }
  nodeInfo.Node = node.Node
  nodeInfo.Devices = policy.DeviceUsageList{
   Policy:      userGPUPolicy,
   DeviceLists: make([]*policy.DeviceListsScore, 0),
  }
  for _, d := range node.Devices {
   nodeInfo.Devices.DeviceLists = append(nodeInfo.Devices.DeviceLists, &policy.DeviceListsScore{
    Score: 0,
    Device: &util.DeviceUsage{
     ID:        d.ID,
     Index:     d.Index,
     Used:      0,
     Count:     d.Count,
     Usedmem:   0,
     Totalmem:  d.Devmem,
     Totalcore: d.Devcore,
     Usedcores: 0,
     MigUsage: util.MigInUse{
      Index:     0,
      UsageList: make(util.MIGS, 0),
     },
     MigTemplate: d.MIGTemplate,
     Mode:        d.Mode,
     Type:        d.Type,
     Numa:        d.Numa,
     Health:      d.Health,
    },
   })
  }
  overallnodeMap[node.ID] = nodeInfo
 }

 podsInfo := s.ListPodsInfo()
 for _, p := range podsInfo {
  node, ok := overallnodeMap[p.NodeID]
  if !ok {
   continue
  }
  for _, podsingleds := range p.Devices {
   for _, ctrdevs := range podsingleds {
    for _, udevice := range ctrdevs {
     for _, d := range node.Devices.DeviceLists {
      deviceID := udevice.UUID
      if strings.Contains(deviceID, "[") {
       deviceID = strings.Split(deviceID, "[")[0]
      }
      if d.Device.ID == deviceID {
       d.Device.Used++
       d.Device.Usedmem += udevice.Usedmem
       d.Device.Usedcores += udevice.Usedcores
       if strings.Contains(udevice.UUID, "[") {
        tmpIdx, Instance := util.ExtractMigTemplatesFromUUID(udevice.UUID)
        if len(d.Device.MigUsage.UsageList) == 0 {
         util.PlatternMIG(&d.Device.MigUsage, d.Device.MigTemplate, tmpIdx)
        }
        d.Device.MigUsage.UsageList[Instance].InUse = true
        klog.V(3).Infoln("add mig usage", d.Device.MigUsage, "template=", d.Device.MigTemplate, "uuid=", d.Device.ID)
       }
      }
     }
    }
   }
  }
  klog.V(5).Infof("usage: pod %v assigned %v %v", p.Name, p.NodeID, p.Devices)
 }
 s.overviewstatus = overallnodeMap
 for _, nodeID := range *nodes {
  node, err := s.GetNode(nodeID)
  if err != nil {
   // The identified node does not have a gpu device, so the log here has no practical meaning,increase log priority.
   klog.V(5).InfoS("node unregistered", "node", nodeID, "error", err)
   failedNodes[nodeID] = "node unregistered"
   continue
  }
  cachenodeMap[node.ID] = overallnodeMap[node.ID]
 }
 s.cachedstatus = cachenodeMap
 return &cachenodeMap, failedNodes, nil
}

获取node总的资源与已分配的资源, 首先获取node信息

pkg/scheduler/nodes.go:120

func (m *nodeManager) ListNodes() (map[string]*util.NodeInfo, error) {
 m.mutex.RLock()
 defer m.mutex.RUnlock()
 return m.nodes, nil
}

这里用到了缓存，缓存节点信息，由 addNode 添加缓存

3.2.3.1.1. Node缓存

pkg/scheduler/nodes.go:46

func (m *nodeManager) addNode(nodeID string, nodeInfo *util.NodeInfo) {
 if nodeInfo == nil || len(nodeInfo.Devices) == 0 {
  return
 }
 m.mutex.Lock()
 defer m.mutex.Unlock()
 _, ok := m.nodes[nodeID]
 if ok {
  if len(nodeInfo.Devices) > 0 {
   tmp := make([]util.DeviceInfo, 0, len(nodeInfo.Devices))
   devices := device.GetDevices()
   deviceType := ""
   for _, val := range devices {
    if strings.Contains(nodeInfo.Devices[0].Type, val.CommonWord()) {
     deviceType = val.CommonWord()
    }
   }
   for _, val := range m.nodes[nodeID].Devices {
    if !strings.Contains(val.Type, deviceType) {
     tmp = append(tmp, val)
    }
   }
   m.nodes[nodeID].Devices = tmp
   m.nodes[nodeID].Devices = append(m.nodes[nodeID].Devices, nodeInfo.Devices...)
  }
 } else {
  m.nodes[nodeID] = nodeInfo
 }
}

这里的主要逻辑在于 device.GetDevices() 获取设备信息

pkg/device/devices.go:81

func GetDevices() map[string]Devices {
 return devices
}

device也是个缓存，后面再分析，首先看Node缓存是什么时候被调用的

pkg/scheduler/scheduler.go:155

func (s *Scheduler) RegisterFromNodeAnnotations() {
 klog.V(5).Infoln("Scheduler into RegisterFromNodeAnnotations")
 ticker := time.NewTicker(time.Second * 15)
 for {
  select {
  case <-s.nodeNotify:
  case <-ticker.C:
  case <-s.stopCh:
   return
  }
  labelSelector := labels.Everything()
  if len(config.NodeLabelSelector) > 0 {
   labelSelector = (labels.Set)(config.NodeLabelSelector).AsSelector()
  }
  rawNodes, err := s.nodeLister.List(labelSelector)
  if err != nil {
   klog.Errorln("nodes list failed", err.Error())
   continue
  }
  var nodeNames []string
  for _, val := range rawNodes {
   nodeNames = append(nodeNames, val.Name)
   for devhandsk, devInstance := range device.GetDevices() {
    health, needUpdate := devInstance.CheckHealth(devhandsk, val)
    klog.V(5).InfoS("device check health", "node", val.Name, "deviceVendor", devhandsk, "health", health, "needUpdate", needUpdate)
    if !health {
     err := devInstance.NodeCleanUp(val.Name)
     // If the device is not healthy, the device is removed from the node.
     // At the same time, this node needs to be removed from the cache.
     if err != nil {
      klog.Errorln("node cleanup failed", err.Error())
     }
     info, ok := s.nodes[val.Name]
     if ok {
      klog.Infof("node %v device %s:%v leave, %v remaining devices:%v", val.Name, devhandsk, info.ID, err, s.nodes[val.Name].Devices)
      s.rmNodeDevice(val.Name, info, devhandsk)
      continue
     }
    }
    if !needUpdate {
     continue
    }
    _, ok := util.HandshakeAnnos[devhandsk]
    if ok {
     tmppat := make(map[string]string)
     tmppat[util.HandshakeAnnos[devhandsk]] = "Requesting_" + time.Now().Format("2006.01.02 15:04:05")
     klog.V(4).InfoS("New timestamp", util.HandshakeAnnos[devhandsk], tmppat[util.HandshakeAnnos[devhandsk]], "nodeName", val.Name)
     n, err := util.GetNode(val.Name)
     if err != nil {
      klog.Errorln("get node failed", err.Error())
      continue
     }
     util.PatchNodeAnnotations(n, tmppat)
    }

    nodeInfo := &util.NodeInfo{}
    nodeInfo.ID = val.Name
    nodeInfo.Node = val
    nodedevices, err := devInstance.GetNodeDevices(*val)
    if err != nil {
     continue
    }
    nodeInfo.Devices = make([]util.DeviceInfo, 0)
    for _, deviceinfo := range nodedevices {
     nodeInfo.Devices = append(nodeInfo.Devices, *deviceinfo)
    }
    s.addNode(val.Name, nodeInfo)
    if s.nodes[val.Name] != nil && len(nodeInfo.Devices) > 0 {
     klog.Infof("node %v device %s come node info=%s,%v total=%v", val.Name, devhandsk, nodeInfo.ID, nodeInfo.Devices, s.nodes[val.Name].Devices)
    }
   }
  }
  _, _, err = s.getNodesUsage(&nodeNames, nil)
  if err != nil {
   klog.Errorln("get node usage failed", err.Error())
  }
 }
}

启动了一个15s的定时任务，获取node信息维护node缓存

这里的核心逻辑在于 for devhandsk, devInstance := range device.GetDevices() 获取所有的device，主要是一些根据不同的设备注册了不同的handler, 根据注册的device获取显卡的资源信息 devInstance.GetNodeDevices

这里会通过注册的device(此环境为nvidia)，调用到不同显卡的GetNodeDevices实现， device后面再做具体说明

pkg/device/nvidia/device.go:209

ffunc (dev *NvidiaGPUDevices) GetNodeDevices(n corev1.Node) ([]*util.DeviceInfo, error) {
 devEncoded, ok := n.Annotations[RegisterAnnos]
 if !ok {
  return []*util.DeviceInfo{}, errors.New("annos not found " + RegisterAnnos)
 }
 nodedevices, err := util.DecodeNodeDevices(devEncoded)
 if err != nil {
  klog.ErrorS(err, "failed to decode node devices", "node", n.Name, "device annotation", devEncoded)
  return []*util.DeviceInfo{}, err
 }
 if len(nodedevices) == 0 {
  klog.InfoS("no nvidia gpu device found", "node", n.Name, "device annotation", devEncoded)
  return []*util.DeviceInfo{}, errors.New("no gpu found on node")
 }
 for _, val := range nodedevices {
  if val.Mode == "mig" {
   val.MIGTemplate = make([]util.Geometry, 0)
   for _, migTemplates := range dev.config.MigGeometriesList {
    found := false
    for _, migDevices := range migTemplates.Models {
     if strings.Contains(val.Type, migDevices) {
      found = true
      break
     }
    }
    if found {
     val.MIGTemplate = append(val.MIGTemplate, migTemplates.Geometries...)
     break
    }
   }
  }
 }
 devDecoded := util.EncodeNodeDevices(nodedevices)
 klog.V(5).InfoS("nodes device information", "node", n.Name, "nodedevices", devDecoded)
 return nodedevices, nil
}

看到这里基本逻辑是 scheduler 通过定时器去读取node的annotation信息并将其维护再node缓存中，以供调度时使用

apiVersion: v1
kind: Node
metadata:
  annotations:
    ...
    hami.io/node-nvidia-register: 'GPU-7aebc545-cbd3-18a0-afce-76cae449702a,10,24576,300,NVIDIA-NVIDIA
      GeForce RTX 3090,0,true:

又调用到了 device，这个我们待会儿再看，继续看谁调用的 RegisterFromNodeAnnotations

cmd/scheduler/main.go:70

func start() {
 device.InitDevices()
 sher = scheduler.NewScheduler()
 sher.Start()
 defer sher.Stop()

 // start monitor metrics
 go sher.RegisterFromNodeAnnotations()
 go initMetrics(config.MetricsBindAddress)

调度器启动的时候就会调用，这里逻辑明确了，继续看刚刚的device

3.2.3.1.2. device

device通过 pkg/device/devices.go:85 进行初始化

func InitDevicesWithConfig(config *Config) {
 devices = make(map[string]Devices)
 DevicesToHandle = []string{}
 devices[nvidia.NvidiaGPUDevice] = nvidia.InitNvidiaDevice(config.NvidiaConfig)
 devices[cambricon.CambriconMLUDevice] = cambricon.InitMLUDevice(config.CambriconConfig)
 devices[hygon.HygonDCUDevice] = hygon.InitDCUDevice(config.HygonConfig)
 devices[iluvatar.IluvatarGPUDevice] = iluvatar.InitIluvatarDevice(config.IluvatarConfig)
 devices[mthreads.MthreadsGPUDevice] = mthreads.InitMthreadsDevice(config.MthreadsConfig)
 devices[metax.MetaxGPUDevice] = metax.InitMetaxDevice(config.MetaxConfig)

 DevicesToHandle = append(DevicesToHandle, nvidia.NvidiaGPUCommonWord)
 DevicesToHandle = append(DevicesToHandle, cambricon.CambriconMLUCommonWord)
 DevicesToHandle = append(DevicesToHandle, hygon.HygonDCUCommonWord)
 DevicesToHandle = append(DevicesToHandle, iluvatar.IluvatarGPUCommonWord)
 DevicesToHandle = append(DevicesToHandle, mthreads.MthreadsGPUCommonWord)
 DevicesToHandle = append(DevicesToHandle, metax.MetaxGPUCommonWord)
 for _, dev := range ascend.InitDevices(config.VNPUs) {
  devices[dev.CommonWord()] = dev
  DevicesToHandle = append(DevicesToHandle, dev.CommonWord())
 }
}

这里用的是nvidia所以主要看 InitNvidiaDevice 即可

pkg/device/devices.go:42

type Devices interface {
 CommonWord() string
 MutateAdmission(ctr *corev1.Container, pod *corev1.Pod) (bool, error)
 CheckHealth(devType string, n *corev1.Node) (bool, bool)
 NodeCleanUp(nn string) error
 GetNodeDevices(n corev1.Node) ([]*util.DeviceInfo, error)
 CheckType(annos map[string]string, d util.DeviceUsage, n util.ContainerDeviceRequest) (bool, bool, bool)
 // CheckUUID is check current device id whether in GPUUseUUID or GPUNoUseUUID set, return true is check success.
 CheckUUID(annos map[string]string, d util.DeviceUsage) bool
 LockNode(n *corev1.Node, p *corev1.Pod) error
 ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) error
 GenerateResourceRequests(ctr *corev1.Container) util.ContainerDeviceRequest
 PatchAnnotations(annoinput *map[string]string, pd util.PodDevices) map[string]string
 CustomFilterRule(allocated *util.PodDevices, request util.ContainerDeviceRequest, toAllicate util.ContainerDevices, device *util.DeviceUsage) bool
 ScoreNode(node *corev1.Node, podDevices util.PodSingleDevice, policy string) float32
 AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error
 // This should not be associated with a specific device object
 //ParseConfig(fs *flag.FlagSet)
}

这里定义了一些接口，不同的设备进行不同的实现，在scheduler启动时进行初始化，以供运行中调用

获取到各个节点的各个设备的资源情况之后开始进行打分

3.2.3.2. 根据节点资源信息打分

pkg/scheduler/scheduler.go:458

 nodeScores, err := s.calcScore(nodeUsage, nums, annos, args.Pod)
 if err != nil {
  err := fmt.Errorf("calcScore failed %v for pod %v", err, args.Pod.Name)
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, err)
  return nil, err
 }

pkg/scheduler/score.go:198

func (s *Scheduler) calcScore(nodes *map[string]*NodeUsage, nums util.PodDeviceRequests, annos map[string]string, task *corev1.Pod) (*policy.NodeScoreList, error) {
 userNodePolicy := config.NodeSchedulerPolicy
 if annos != nil {
  if value, ok := annos[policy.NodeSchedulerPolicyAnnotationKey]; ok {
   userNodePolicy = value
  }
 }
 res := policy.NodeScoreList{
  Policy:   userNodePolicy,
  NodeList: make([]*policy.NodeScore, 0),
 }

 //func calcScore(nodes *map[string]*NodeUsage, errMap *map[string]string, nums util.PodDeviceRequests, annos map[string]string, task *corev1.Pod) (*NodeScoreList, error) {
 // res := make(NodeScoreList, 0, len(*nodes))
 for nodeID, node := range *nodes {
  viewStatus(*node)
  score := policy.NodeScore{NodeID: nodeID, Node: node.Node, Devices: make(util.PodDevices), Score: 0}
  score.ComputeDefaultScore(node.Devices)

  //This loop is for different container request
  ctrfit := false
  for ctrid, n := range nums {
   sums := 0
   for _, k := range n {
    sums += int(k.Nums)
   }

   if sums == 0 {
    for idx := range score.Devices {
     for len(score.Devices[idx]) <= ctrid {
      score.Devices[idx] = append(score.Devices[idx], util.ContainerDevices{})
     }
     score.Devices[idx][ctrid] = append(score.Devices[idx][ctrid], util.ContainerDevice{})
     continue
    }
   }
   klog.V(5).InfoS("fitInDevices", "pod", klog.KObj(task), "node", nodeID)
   fit, _ := fitInDevices(node, n, annos, task, &score.Devices)
   ctrfit = fit
   if !fit {
    klog.InfoS("calcScore:node not fit pod", "pod", klog.KObj(task), "node", nodeID)
    break
   }
  }

  if ctrfit {
   res.NodeList = append(res.NodeList, &score)
   score.OverrideScore(node.Devices, userNodePolicy)
  }
 }
 return &res, nil
}

这块逻辑主要分为遍历节点打分，遍历pod的容器计算每个容器对应的设备的分数，返回所有可以承载limits所需资源的node返回

3.2.3.3. 计算出节点的分数

pkg/scheduler/policy/node_policy.go:68

func (ns *NodeScore) ComputeDefaultScore(devices DeviceUsageList) {
 used, usedCore, usedMem := int32(0), int32(0), int32(0)
 for _, device := range devices.DeviceLists {
  used += device.Device.Used
  usedCore += device.Device.Usedcores
  usedMem += device.Device.Usedmem
 }
 klog.V(2).Infof("node %s used %d, usedCore %d, usedMem %d,", ns.NodeID, used, usedCore, usedMem)

 total, totalCore, totalMem := int32(0), int32(0), int32(0)
 for _, deviceLists := range devices.DeviceLists {
  total += deviceLists.Device.Count
  totalCore += deviceLists.Device.Totalcore
  totalMem += deviceLists.Device.Totalmem
 }
 useScore := float32(used) / float32(total)
 coreScore := float32(usedCore) / float32(totalCore)
 memScore := float32(usedMem) / float32(totalMem)
 ns.Score = float32(Weight) * (useScore + coreScore + memScore)
 klog.V(2).Infof("node %s computer default score is %f", ns.NodeID, ns.Score)
}

节点打分规则比较简单

3.2.3.4. 计算每个容器对应的设备的分数

pkg/scheduler/score.go:149

func fitInDevices(node *NodeUsage, requests util.ContainerDeviceRequests, annos map[string]string, pod *corev1.Pod, devinput *util.PodDevices) (bool, float32) {
 //devmap := make(map[string]util.ContainerDevices)
 devs := util.ContainerDevices{}
 total, totalCore, totalMem := int32(0), int32(0), int32(0)
 free, freeCore, freeMem := int32(0), int32(0), int32(0)
 sums := 0
 // computer all device score for one node
 for index := range node.Devices.DeviceLists {
  node.Devices.DeviceLists[index].ComputeScore(requests)
 }
 //This loop is for requests for different devices
 for _, k := range requests {
  sums += int(k.Nums)
  if int(k.Nums) > len(node.Devices.DeviceLists) {
   klog.InfoS("request devices nums cannot exceed the total number of devices on the node.", "pod", klog.KObj(pod), "request devices nums", k.Nums, "node device nums", len(node.Devices.DeviceLists))
   return false, 0
  }
  sort.Sort(node.Devices)
  fit, tmpDevs := fitInCertainDevice(node, k, annos, pod, devinput)
  if fit {
   for idx, val := range tmpDevs[k.Type] {
    for nidx, v := range node.Devices.DeviceLists {
     //bc node.Devices has been sorted, so we should find out the correct device
     if v.Device.ID != val.UUID {
      continue
     }
     total += v.Device.Count
     totalCore += v.Device.Totalcore
     totalMem += v.Device.Totalmem
     free += v.Device.Count - v.Device.Used
     freeCore += v.Device.Totalcore - v.Device.Usedcores
     freeMem += v.Device.Totalmem - v.Device.Usedmem
     err := device.GetDevices()[k.Type].AddResourceUsage(node.Devices.DeviceLists[nidx].Device, &tmpDevs[k.Type][idx])
     if err != nil {
      klog.Errorf("AddResource failed:%s", err.Error())
      return false, 0
     }
     klog.Infoln("After AddResourceUsage:", node.Devices.DeviceLists[nidx].Device)
    }
   }
   devs = append(devs, tmpDevs[k.Type]...)
  } else {
   return false, 0
  }
  (*devinput)[k.Type] = append((*devinput)[k.Type], devs)
 }
 return true, 0
}

主要逻辑为

给容器对应的每个设备打分、遍历不同的容器对应的limit资源，找到可以承载容器limits资源的设备

pkg/scheduler/policy/gpu_policy.go:58

func (ds *DeviceListsScore) ComputeScore(requests util.ContainerDeviceRequests) {
 request, core, mem := int32(0), int32(0), int32(0)
 // Here we are required to use the same type device
 for _, container := range requests {
  request += container.Nums
  core += container.Coresreq
  if container.MemPercentagereq != 0 && container.MemPercentagereq != 101 {
   mem += ds.Device.Totalmem * (container.MemPercentagereq / 100.0)
   continue
  }
  mem += container.Memreq
 }
 klog.V(2).Infof("device %s user %d, userCore %d, userMem %d,", ds.Device.ID, ds.Device.Used, ds.Device.Usedcores, ds.Device.Usedmem)

 usedScore := float32(request+ds.Device.Used) / float32(ds.Device.Count)
 coreScore := float32(core+ds.Device.Usedcores) / float32(ds.Device.Totalcore)
 memScore := float32(mem+ds.Device.Usedmem) / float32(ds.Device.Totalmem)
 ds.Score = float32(Weight) * (usedScore + coreScore + memScore)
 klog.V(2).Infof("device %s computer score is %f", ds.Device.ID, ds.Score)
}

打分规则与节点类似

pkg/scheduler/score.go:65

func fitInCertainDevice(node *NodeUsage, request util.ContainerDeviceRequest, annos map[string]string, pod *corev1.Pod, allocated *util.PodDevices) (bool, map[string]util.ContainerDevices) {
 k := request
 originReq := k.Nums
 prevnuma := -1
 klog.InfoS("Allocating device for container request", "pod", klog.KObj(pod), "card request", k)
 var tmpDevs map[string]util.ContainerDevices
 tmpDevs = make(map[string]util.ContainerDevices)
 for i := len(node.Devices.DeviceLists) - 1; i >= 0; i-- {
  klog.InfoS("scoring pod", "pod", klog.KObj(pod), "Memreq", k.Memreq, "MemPercentagereq", k.MemPercentagereq, "Coresreq", k.Coresreq, "Nums", k.Nums, "device index", i, "device", node.Devices.DeviceLists[i].Device.ID)
  found, numa := checkType(annos, *node.Devices.DeviceLists[i].Device, k)
  if !found {
   klog.InfoS("card type mismatch,continuing...", "pod", klog.KObj(pod), (node.Devices.DeviceLists[i].Device).Type, k.Type)
   continue
  }
  if numa && prevnuma != node.Devices.DeviceLists[i].Device.Numa {
   klog.InfoS("Numa not fit, resotoreing", "pod", klog.KObj(pod), "k.nums", k.Nums, "numa", numa, "prevnuma", prevnuma, "device numa", node.Devices.DeviceLists[i].Device.Numa)
   k.Nums = originReq
   prevnuma = node.Devices.DeviceLists[i].Device.Numa
   tmpDevs = make(map[string]util.ContainerDevices)
  }
  if !checkUUID(annos, *node.Devices.DeviceLists[i].Device, k) {
   klog.InfoS("card uuid mismatch,", "pod", klog.KObj(pod), "current device info is:", *node.Devices.DeviceLists[i].Device)
   continue
  }

  memreq := int32(0)
  if node.Devices.DeviceLists[i].Device.Count <= node.Devices.DeviceLists[i].Device.Used {
   continue
  }
  if k.Coresreq > 100 {
   klog.ErrorS(nil, "core limit can't exceed 100", "pod", klog.KObj(pod))
   k.Coresreq = 100
   //return false, tmpDevs
  }
  if k.Memreq > 0 {
   memreq = k.Memreq
  }
  if k.MemPercentagereq != 101 && k.Memreq == 0 {
   //This incurs an issue
   memreq = node.Devices.DeviceLists[i].Device.Totalmem * k.MemPercentagereq / 100
  }
  if node.Devices.DeviceLists[i].Device.Totalmem-node.Devices.DeviceLists[i].Device.Usedmem < memreq {
   klog.V(5).InfoS("card Insufficient remaining memory", "pod", klog.KObj(pod), "device index", i, "device", node.Devices.DeviceLists[i].Device.ID, "device total memory", node.Devices.DeviceLists[i].Device.Totalmem, "device used memory", node.Devices.DeviceLists[i].Device.Usedmem, "request memory", memreq)
   continue
  }
  if node.Devices.DeviceLists[i].Device.Totalcore-node.Devices.DeviceLists[i].Device.Usedcores < k.Coresreq {
   klog.V(5).InfoS("card Insufficient remaining cores", "pod", klog.KObj(pod), "device index", i, "device", node.Devices.DeviceLists[i].Device.ID, "device total core", node.Devices.DeviceLists[i].Device.Totalcore, "device used core", node.Devices.DeviceLists[i].Device.Usedcores, "request cores", k.Coresreq)
   continue
  }
  // Coresreq=100 indicates it want this card exclusively
  if node.Devices.DeviceLists[i].Device.Totalcore == 100 && k.Coresreq == 100 && node.Devices.DeviceLists[i].Device.Used > 0 {
   klog.V(5).InfoS("the container wants exclusive access to an entire card, but the card is already in use", "pod", klog.KObj(pod), "device index", i, "device", node.Devices.DeviceLists[i].Device.ID, "used", node.Devices.DeviceLists[i].Device.Used)
   continue
  }
  // You can't allocate core=0 job to an already full GPU
  if node.Devices.DeviceLists[i].Device.Totalcore != 0 && node.Devices.DeviceLists[i].Device.Usedcores == node.Devices.DeviceLists[i].Device.Totalcore && k.Coresreq == 0 {
   klog.V(5).InfoS("can't allocate core=0 job to an already full GPU", "pod", klog.KObj(pod), "device index", i, "device", node.Devices.DeviceLists[i].Device.ID)
   continue
  }
  if !device.GetDevices()[k.Type].CustomFilterRule(allocated, request, tmpDevs[k.Type], node.Devices.DeviceLists[i].Device) {
   continue
  }
  if k.Nums > 0 {
   klog.InfoS("first fitted", "pod", klog.KObj(pod), "device", node.Devices.DeviceLists[i].Device.ID)
   k.Nums--
   tmpDevs[k.Type] = append(tmpDevs[k.Type], util.ContainerDevice{
    Idx:       int(node.Devices.DeviceLists[i].Device.Index),
    UUID:      node.Devices.DeviceLists[i].Device.ID,
    Type:      k.Type,
    Usedmem:   memreq,
    Usedcores: k.Coresreq,
   })
  }
  if k.Nums == 0 {
   klog.InfoS("device allocate success", "pod", klog.KObj(pod), "allocate device", tmpDevs)
   return true, tmpDevs
  }
  if node.Devices.DeviceLists[i].Device.Mode == "mig" {
   i++
  }
 }
 return false, tmpDevs
}

遍历设备，主要根据设备资源余量来判断是否够container分配，返回所有够分配的设备

pkg/scheduler/scheduler.go:458

 nodeScores, err := s.calcScore(nodeUsage, nums, annos, args.Pod)
 if err != nil {
  err := fmt.Errorf("calcScore failed %v for pod %v", err, args.Pod.Name)
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, err)
  return nil, err
 }
 if len((*nodeScores).NodeList) == 0 {
  klog.V(4).Infof("All node scores do not meet for pod %v", args.Pod.Name)
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, fmt.Errorf("no available node, all node scores do not meet"))
  return &extenderv1.ExtenderFilterResult{
   FailedNodes: failedNodes,
  }, nil
 }
 klog.V(4).Infoln("nodeScores_len=", len((*nodeScores).NodeList))
 sort.Sort(nodeScores)
 m := (*nodeScores).NodeList[len((*nodeScores).NodeList)-1]
 klog.Infof("schedule %v/%v to %v %v", args.Pod.Namespace, args.Pod.Name, m.NodeID, m.Devices)
 annotations := make(map[string]string)
 annotations[util.AssignedNodeAnnotations] = m.NodeID
 annotations[util.AssignedTimeAnnotations] = strconv.FormatInt(time.Now().Unix(), 10)

  for _, val := range device.GetDevices() {
  val.PatchAnnotations(&annotations, m.Devices)
 }

 //InRequestDevices := util.EncodePodDevices(util.InRequestDevices, m.devices)
 //supportDevices := util.EncodePodDevices(util.SupportDevices, m.devices)
 //maps.Copy(annotations, InRequestDevices)
 //maps.Copy(annotations, supportDevices)
 s.addPod(args.Pod, m.NodeID, m.Devices)
 err = util.PatchPodAnnotations(args.Pod, annotations)
 if err != nil {
  s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringFailed, []string{}, err)
  s.delPod(args.Pod)
  return nil, err
 }
 s.recordScheduleFilterResultEvent(args.Pod, EventReasonFilteringSucceed, []string{m.NodeID}, nil)
 res := extenderv1.ExtenderFilterResult{NodeNames: &[]string{m.NodeID}}
 return &res, nil

遍历完成之后选择分数最高的, 给pod打标签

apiVersion: v1
kind: Pod
metadata:
  annotations:
    hami.io/vgpu-node: node1
    hami.io/vgpu-time: "1733988480"
    hami.io/vgpu-devices-allocated: GPU-7aebc545-cbd3-18a0-afce-76cae449702a,NVIDIA,20000,80:;
    hami.io/vgpu-devices-to-allocate: ;

3.2.4. binding实现

bind逻辑比较简单，将pod绑定到node

pkg/scheduler/routes/route.go:82

func Bind(s *scheduler.Scheduler) httprouter.Handle {
 return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
  var buf bytes.Buffer
  body := io.TeeReader(r.Body, &buf)
  var extenderBindingArgs extenderv1.ExtenderBindingArgs
  var extenderBindingResult *extenderv1.ExtenderBindingResult

  if err := json.NewDecoder(body).Decode(&extenderBindingArgs); err != nil {
   klog.ErrorS(err, "Decode extender binding args")
   extenderBindingResult = &extenderv1.ExtenderBindingResult{
    Error: err.Error(),
   }
  } else {
   extenderBindingResult, err = s.Bind(extenderBindingArgs)
  }

  if response, err := json.Marshal(extenderBindingResult); err != nil {
   klog.ErrorS(err, "Marshal binding result", "result", extenderBindingResult)
   w.Header().Set("Content-Type", "application/json")
   w.WriteHeader(http.StatusInternalServerError)
   errMsg := fmt.Sprintf("{'error':'%s'}", err.Error())
   w.Write([]byte(errMsg))
  } else {
   klog.V(5).InfoS("Return bind response", "result", extenderBindingResult)
   w.Header().Set("Content-Type", "application/json")
   w.WriteHeader(http.StatusOK)
   w.Write(response)
  }
 }
}

路由处理

func (s *Scheduler) Bind(args extenderv1.ExtenderBindingArgs) (*extenderv1.ExtenderBindingResult, error) {
 klog.InfoS("Bind", "pod", args.PodName, "namespace", args.PodNamespace, "podUID", args.PodUID, "node", args.Node)
 var err error
 var res *extenderv1.ExtenderBindingResult
 binding := &corev1.Binding{
  ObjectMeta: metav1.ObjectMeta{Name: args.PodName, UID: args.PodUID},
  Target:     corev1.ObjectReference{Kind: "Node", Name: args.Node},
 }
 current, err := s.kubeClient.CoreV1().Pods(args.PodNamespace).Get(context.Background(), args.PodName, metav1.GetOptions{})
 if err != nil {
  klog.ErrorS(err, "Get pod failed")
 }

 node, err := s.kubeClient.CoreV1().Nodes().Get(context.Background(), args.Node, metav1.GetOptions{})
 if err != nil {
  klog.ErrorS(err, "Failed to get node", "node", args.Node)
  s.recordScheduleBindingResultEvent(current, EventReasonBindingFailed, []string{}, fmt.Errorf("failed to get node %v", args.Node))
  res = &extenderv1.ExtenderBindingResult{
   Error: err.Error(),
  }
  return res, nil
 }

 tmppatch := make(map[string]string)
 for _, val := range device.GetDevices() {
  err = val.LockNode(node, current)
  if err != nil {
   goto ReleaseNodeLocks
  }
 }

 tmppatch[util.DeviceBindPhase] = "allocating"
 tmppatch[util.BindTimeAnnotations] = strconv.FormatInt(time.Now().Unix(), 10)

 err = util.PatchPodAnnotations(current, tmppatch)
 if err != nil {
  klog.ErrorS(err, "patch pod annotation failed")
 }
 if err = s.kubeClient.CoreV1().Pods(args.PodNamespace).Bind(context.Background(), binding, metav1.CreateOptions{}); err != nil {
  klog.ErrorS(err, "Failed to bind pod", "pod", args.PodName, "namespace", args.PodNamespace, "podUID", args.PodUID, "node", args.Node)
 }
 if err == nil {
  s.recordScheduleBindingResultEvent(current, EventReasonBindingSucceed, []string{args.Node}, nil)
  res = &extenderv1.ExtenderBindingResult{
   Error: "",
  }
  klog.Infoln("After Binding Process")
  return res, nil
 }
ReleaseNodeLocks:
 klog.InfoS("bind failed", "err", err.Error())
 for _, val := range device.GetDevices() {
  val.ReleaseNodeLock(node, current)
 }
 s.recordScheduleBindingResultEvent(current, EventReasonBindingFailed, []string{}, err)
 return &extenderv1.ExtenderBindingResult{
  Error: err.Error(),
 }, nil
}

3.3. Node将设备情况写入node annotation

scheduler获取node的设备信息主要是通过读取node的annotation，主要有如下几步

启动插件

apiVersion: v1
kind: Node
metadata:
  annotations:
    hami.io/node-handshake: Requesting_2024.12.24 03:31:30
    hami.io/node-handshake-dcu: Deleted_2024.12.06 07:43:49
    hami.io/node-nvidia-register: 'GPU-7aebc545-cbd3-18a0-afce-76cae449702a,10,73728,300,NVIDIA-NVIDIA
      GeForce RTX 3090,0,true:'

3.3.1. 启动device-plugin服务

这里用到了 github.com/urfave/cli/v2 作为command启动服务，需要注意 -v不是日志等级而是是否显示版本

cmd/device-plugin/nvidia/main.go:40

func main() {
 var configFile string

 c := cli.NewApp()
 c.Name = "NVIDIA Device Plugin"
 c.Usage = "NVIDIA device plugin for Kubernetes"
 c.Version = info.GetVersionString()
 c.Action = func(ctx *cli.Context) error {
  return start(ctx, c.Flags)
 }

3.3.2. 启动plugin

这里的plugin主要是针对不同厂家的设备需要实现不同的方法，这里定义了pluigin的控制器，例如start、restart、exit等，这里我们主要关注plugins, restartPlugins, err := startPlugins(c, flags, restarting)

cmd/device-plugin/nvidia/main.go:156

func start(c *cli.Context, flags []cli.Flag) error {
 klog.Info("Starting FS watcher.")
 util.NodeName = os.Getenv(util.NodeNameEnvName)
 watcher, err := newFSWatcher(kubeletdevicepluginv1beta1.DevicePluginPath)
 if err != nil {
  return fmt.Errorf("failed to create FS watcher: %v", err)
 }
 defer watcher.Close()
 //device.InitDevices()

 /*Loading config files*/
 klog.Infof("Start working on node %s", util.NodeName)
 klog.Info("Starting OS watcher.")
 sigs := newOSWatcher(syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

 var restarting bool
 var restartTimeout <-chan time.Time
 var plugins []plugin.Interface
restart:
 // If we are restarting, stop plugins from previous run.
 if restarting {
  err := stopPlugins(plugins)
  if err != nil {
   return fmt.Errorf("error stopping plugins from previous run: %v", err)
  }
 }

 klog.Info("Starting Plugins.")
 plugins, restartPlugins, err := startPlugins(c, flags, restarting)
 if err != nil {
  return fmt.Errorf("error starting plugins: %v", err)
 }

 if restartPlugins {
  klog.Info("Failed to start one or more plugins. Retrying in 30s...")
  restartTimeout = time.After(30 * time.Second)
 }

 restarting = true

 // Start an infinite loop, waiting for several indicators to either log
 // some messages, trigger a restart of the plugins, or exit the program.
 for {
  select {
  // If the restart timeout has expired, then restart the plugins
  case <-restartTimeout:
   goto restart

  // Detect a kubelet restart by watching for a newly created
  // 'kubeletdevicepluginv1beta1.KubeletSocket' file. When this occurs, restart this loop,
  // restarting all of the plugins in the process.
  case event := <-watcher.Events:
   if event.Name == kubeletdevicepluginv1beta1.KubeletSocket && event.Op&fsnotify.Create == fsnotify.Create {
    klog.Infof("inotify: %s created, restarting.", kubeletdevicepluginv1beta1.KubeletSocket)
    goto restart
   }

  // Watch for any other fs errors and log them.
  case err := <-watcher.Errors:
   klog.Errorf("inotify: %s", err)

  // Watch for any signals from the OS. On SIGHUP, restart this loop,
  // restarting all of the plugins in the process. On all other
  // signals, exit the loop and exit the program.
  case s := <-sigs:
   switch s {
   case syscall.SIGHUP:
    klog.Info("Received SIGHUP, restarting.")
    goto restart
   default:
    klog.Infof("Received signal \"%v\", shutting down.", s)
    goto exit
   }
  }
 }
exit:
 err = stopPlugins(plugins)
 if err != nil {
  return fmt.Errorf("error stopping plugins: %v", err)
 }
 return nil
}

cmd/device-plugin/nvidia/main.go:239

启动插件，主要方法 p.Start()

func startPlugins(c *cli.Context, flags []cli.Flag, restarting bool) ([]plugin.Interface, bool, error) {
 // Load the configuration file
 klog.Info("Loading configuration.")
 config, err := loadConfig(c, flags)
 if err != nil {
  return nil, false, fmt.Errorf("unable to load config: %v", err)
 }
 disableResourceRenamingInConfig(config)

 /*Loading config files*/
 //fmt.Println("NodeName=", config.NodeName)
 devConfig, err := generateDeviceConfigFromNvidia(config, c, flags)
 if err != nil {
  klog.Errorf("failed to load config file %s", err.Error())
  return nil, false, err
 }

 // Update the configuration file with default resources.
 klog.Info("Updating config with default resource matching patterns.")
 err = rm.AddDefaultResourcesToConfig(&devConfig)
 if err != nil {
  return nil, false, fmt.Errorf("unable to add default resources to config: %v", err)
 }

 // Print the config to the output.
 configJSON, err := json.MarshalIndent(devConfig, "", "  ")
 if err != nil {
  return nil, false, fmt.Errorf("failed to marshal config to JSON: %v", err)
 }
 klog.Infof("\nRunning with config:\n%v", string(configJSON))

 // Get the set of plugins.
 klog.Info("Retrieving plugins.")
 pluginManager, err := NewPluginManager(&devConfig)
 if err != nil {
  return nil, false, fmt.Errorf("error creating plugin manager: %v", err)
 }
 plugins, err := pluginManager.GetPlugins()
 if err != nil {
  return nil, false, fmt.Errorf("error getting plugins: %v", err)
 }

 // Loop through all plugins, starting them if they have any devices
 // to serve. If even one plugin fails to start properly, try
 // starting them all again.
 started := 0
 for _, p := range plugins {
  // Just continue if there are no devices to serve for plugin p.
  if len(p.Devices()) == 0 {
   continue
  }

  // Start the gRPC server for plugin p and connect it with the kubelet.
  if err := p.Start(); err != nil {
   klog.Error("Could not contact Kubelet. Did you enable the device plugin feature gate?")
   klog.Error("You can check the prerequisites at: https://github.com/NVIDIA/k8s-device-plugin#prerequisites")
   klog.Error("You can learn how to set the runtime at: https://github.com/NVIDIA/k8s-device-plugin#quick-start")
   return plugins, true, nil
  }
  started++
 }

 if started == 0 {
  klog.Info("No devices found. Waiting indefinitely.")
 }

 return plugins, false, nil
}

其中p(plugin)需要实现几个方法来管理插件

pkg/device-plugin/nvidiadevice/nvinternal/plugin/api.go:37

type Interface interface {
 Devices() rm.Devices
 Start() error
 Stop() error
}

同时如果需要kubelet能够识别 resource中的类似 nvidia.com/gpu: 1 这样的拓展字段需要启动一个grpc服务挂载 /var/lib/kubelet/device-plugins/ 且实现如下方法，这块跟调度相关性不大，暂且不展开 device-plugins

k8s.io/kubelet@v0.28.3/pkg/apis/deviceplugin/v1beta1/api.pb.go:1419

type DevicePluginServer interface {
 // GetDevicePluginOptions returns options to be communicated with Device
 // Manager
 GetDevicePluginOptions(context.Context, *Empty) (*DevicePluginOptions, error)
 // ListAndWatch returns a stream of List of Devices
 // Whenever a Device state change or a Device disappears, ListAndWatch
 // returns the new list
 ListAndWatch(*Empty, DevicePlugin_ListAndWatchServer) error
 // GetPreferredAllocation returns a preferred set of devices to allocate
 // from a list of available ones. The resulting preferred allocation is not
 // guaranteed to be the allocation ultimately performed by the
 // devicemanager. It is only designed to help the devicemanager make a more
 // informed allocation decision when possible.
 GetPreferredAllocation(context.Context, *PreferredAllocationRequest) (*PreferredAllocationResponse, error)
 // Allocate is called during container creation so that the Device
 // Plugin can run device specific operations and instruct Kubelet
 // of the steps to make the Device available in the container
 Allocate(context.Context, *AllocateRequest) (*AllocateResponse, error)
 // PreStartContainer is called, if indicated by Device Plugin during registeration phase,
 // before each container start. Device plugin can run device specific operations
 // such as resetting the device before making devices available to the container
 PreStartContainer(context.Context, *PreStartContainerRequest) (*PreStartContainerResponse, error)
}

3.3.3. nvidia插件的实现

主要看plugin.WatchAndRegister()

pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go:196

func (plugin *NvidiaDevicePlugin) Start() error {
 plugin.initialize()

 err := plugin.Serve()
 if err != nil {
  klog.Infof("Could not start device plugin for '%s': %s", plugin.rm.Resource(), err)
  plugin.cleanup()
  return err
 }
 klog.Infof("Starting to serve '%s' on %s", plugin.rm.Resource(), plugin.socket)

 err = plugin.Register()
 if err != nil {
  klog.Infof("Could not register device plugin: %s", err)
  plugin.Stop()
  return err
 }
 klog.Infof("Registered device plugin for '%s' with Kubelet", plugin.rm.Resource())

 if plugin.operatingMode == "mig" {
  cmd := exec.Command("nvidia-mig-parted", "export")
  var stdout, stderr bytes.Buffer
  cmd.Stdout = &stdout
  cmd.Stderr = &stderr
  err := cmd.Run()
  if err != nil {
   klog.Fatalf("nvidia-mig-parted failed with %s\n", err)
  }
  outStr := stdout.Bytes()
  yaml.Unmarshal(outStr, &plugin.migCurrent)
  os.WriteFile("/tmp/migconfig.yaml", outStr, os.ModePerm)
  if len(plugin.migCurrent.MigConfigs["current"]) == 1 && len(plugin.migCurrent.MigConfigs["current"][0].Devices) == 0 {
   idx := 0
   plugin.migCurrent.MigConfigs["current"][0].Devices = make([]int32, 0)
   for idx < GetDeviceNums() {
    plugin.migCurrent.MigConfigs["current"][0].Devices = append(plugin.migCurrent.MigConfigs["current"][0].Devices, int32(idx))
    idx++
   }
  }
  klog.Infoln("Mig export", plugin.migCurrent)
 }
 go func() {
  err := plugin.rm.CheckHealth(plugin.stop, plugin.health)
  if err != nil {
   klog.Infof("Failed to start health check: %v; continuing with health checks disabled", err)
  }
 }()

 go func() {
  plugin.WatchAndRegister()
 }()

 return nil
}

这里是个定时器，每30s收集一次该node的设备信息，并写入node annotation

func (plugin *NvidiaDevicePlugin) WatchAndRegister() {
 klog.Info("Starting WatchAndRegister")
 errorSleepInterval := time.Second * 5
 successSleepInterval := time.Second * 30
 for {
  err := plugin.RegistrInAnnotation()
  if err != nil {
   klog.Errorf("Failed to register annotation: %v", err)
   klog.Infof("Retrying in %v seconds...", errorSleepInterval)
   time.Sleep(errorSleepInterval)
  } else {
   klog.Infof("Successfully registered annotation. Next check in %v seconds...", successSleepInterval)
   time.Sleep(successSleepInterval)
  }
 }
}

func (plugin *NvidiaDevicePlugin) RegistrInAnnotation() error {
 devices := plugin.getAPIDevices()
 klog.InfoS("start working on the devices", "devices", devices)
 annos := make(map[string]string)
 node, err := util.GetNode(util.NodeName)
 if err != nil {
  klog.Errorln("get node error", err.Error())
  return err
 }
 encodeddevices := util.EncodeNodeDevices(*devices)
 annos[nvidia.HandshakeAnnos] = "Reported " + time.Now().String()
 annos[nvidia.RegisterAnnos] = encodeddevices
 klog.Infof("patch node with the following annos %v", fmt.Sprintf("%v", annos))
 err = util.PatchNodeAnnotations(node, annos)

 if err != nil {
  klog.Errorln("patch node error", err.Error())
 }
 return err
}

具体数据收集逻辑

pkg/device-plugin/nvidiadevice/nvinternal/plugin/register.go:110

func (plugin *NvidiaDevicePlugin) getAPIDevices() *[]*util.DeviceInfo {
 devs := plugin.Devices()
 klog.V(5).InfoS("getAPIDevices", "devices", devs)
 nvml.Init()
 res := make([]*util.DeviceInfo, 0, len(devs))
 for UUID := range devs {
  ndev, ret := nvml.DeviceGetHandleByUUID(UUID)
  if ret != nvml.SUCCESS {
   klog.Errorln("nvml new device by index error uuid=", UUID, "err=", ret)
   panic(0)
  }
  idx, ret := ndev.GetIndex()
  if ret != nvml.SUCCESS {
   klog.Errorln("nvml get index error ret=", ret)
   panic(0)
  }
  memoryTotal := 0
  memory, ret := ndev.GetMemoryInfo()
  if ret == nvml.SUCCESS {
   memoryTotal = int(memory.Total)
  } else {
   klog.Error("nvml get memory error ret=", ret)
   panic(0)
  }
  Model, ret := ndev.GetName()
  if ret != nvml.SUCCESS {
   klog.Error("nvml get name error ret=", ret)
   panic(0)
  }

  registeredmem := int32(memoryTotal / 1024 / 1024)
  if plugin.schedulerConfig.DeviceMemoryScaling != 1 {
   registeredmem = int32(float64(registeredmem) * plugin.schedulerConfig.DeviceMemoryScaling)
  }
  klog.Infoln("MemoryScaling=", plugin.schedulerConfig.DeviceMemoryScaling, "registeredmem=", registeredmem)
  health := true
  for _, val := range devs {
   if strings.Compare(val.ID, UUID) == 0 {
    // when NVIDIA-Tesla P4, the device info is : ID:GPU-e290caca-2f0c-9582-acab-67a142b61ffa,Health:Healthy,Topology:nil,
    // it is more reasonable to think of healthy as case-insensitive
    if strings.EqualFold(val.Health, "healthy") {
     health = true
    } else {
     health = false
    }
    break
   }
  }
  numa, err := plugin.getNumaInformation(idx)
  if err != nil {
   klog.ErrorS(err, "failed to get numa information", "idx", idx)
  }
  res = append(res, &util.DeviceInfo{
   ID:      UUID,
   Index:   uint(idx),
   Count:   int32(plugin.schedulerConfig.DeviceSplitCount),
   Devmem:  registeredmem,
   Devcore: int32(plugin.schedulerConfig.DeviceCoreScaling * 100),
   Type:    fmt.Sprintf("%v-%v", "NVIDIA", Model),
   Numa:    numa,
   Mode:    plugin.operatingMode,
   Health:  health,
  })
  klog.Infof("nvml registered device id=%v, memory=%v, type=%v, numa=%v", idx, registeredmem, Model, numa)
 }
 return &res
}

这里通过nvidia驱动获取设备信息，需要注意的是这里有配置DeviceMemoryScaling, 内存超分配置，这里是通过命令行启动的–config-file 参数指定的schduler配置和代码中固化的config/config.json 来取值的，其中config/config.json 优先级大于–config-file

到这里，调度所需的所有东西就准备好了，pod可以顺利被分配到合适的节点上

4. 参考

HAMI项目本地调试

Fri, 20 Dec 2024 00:00:00 +0000

HAMI项目本地调试

English version generated by AI. 查看英文版

HAMI项目本地调试

1. 前提

在k8s集群上已经有一套安装好并且功能正常的HAMI集群，这样在调试的时候有很多配置可以复用。

hami版本: v2.4.1
OS: Ubuntu 22.04.4 LTS
kernel: 5.15.0-125-generic
nvcc: V11.5.119
nvidia-smi: 550.54.14
CUDA: 12.4

nvidia-smi
Fri Dec 20 02:45:30 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.14              Driver Version: 550.54.14      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3090        Off |   00000000:00:06.0 Off |                  N/A |
|  0%   26C    P8             11W /  370W |       0MiB /  24576MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

2. 创建本地目录保存本地配置

本地启动需要依赖一些文件，可使用统一的文件夹存储

2.1. 创建文件夹

例如，在项目下创建 config 目录

mkdir config

3. hami-device本地调试

3.1. 收集配置

hami-device需要用到两个配置文件，分别为

config/device-config.yaml
config.json

获取并保存到到本地目录

k exec  hami-device-plugin-fn8lc -c device-plugin -- cat /device-config.yaml > /config/device.yaml
k exec hami-device-plugin-fn8lc -c device-plugin -- cat /config/config.json > /config/config.json

3.2. 收集环境变量

这里有个小技巧，在导出pod环境变量前，为pod配置 enableServiceLinks, 例如deployment配置 spec.template.spec.enableServiceLinks: false, 即可去除k8s自动加上去的信息，仅保留pod本身配置的环境变量

k exec hami-device-plugin-lzm4n env >  start-plugin.sh

注意⚠️: 导出的文件里有个PATH变量，这里需要改成 export ATH=$PATH:xxxxxx xxx用原来的输出替代; 每行前面加上export 使用dlv配合goland进行debug

删除带有IP的配置;保留端口配置;增加debug信息, 最终得到一个如下文件

#!/usr/bin/bash
export PATH=$PATH/k8s-vgpu/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
export HOSTNAME=hami-scheduler-858f6b9dcf-4nn58
export KUBERNETES_SERVICE_PORT=443
export KUBERNETES_SERVICE_PORT_HTTPS=443
export KUBERNETES_PORT=tcp://10.233.0.1:443
export KUBERNETES_PORT_443_TCP=tcp://10.233.0.1:443
export KUBERNETES_PORT_443_TCP_PROTO=tcp
export KUBERNETES_PORT_443_TCP_PORT=443
export KUBERNETES_PORT_443_TCP_ADDR=10.233.0.1
export KUBERNETES_SERVICE_HOST=10.233.0.1
export NVARCH=x86_64
export NVIDIA_REQUIRE_CUDA=cuda>=12.6 brand=unknown,driver>=470,driver<471 brand=grid,driver>=470,driver<471 brand=tesla,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=quadro,driver>=470,driver<471 brand=quadrortx,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471 brand=vapps,driver>=470,driver<471 brand=vpc,driver>=470,driver<471 brand=vcs,driver>=470,driver<471 brand=vws,driver>=470,driver<471 brand=cloudgaming,driver>=470,driver<471 brand=unknown,driver>=535,driver<536 brand=grid,driver>=535,driver<536 brand=tesla,driver>=535,driver<536 brand=nvidia,driver>=535,driver<536 brand=quadro,driver>=535,driver<536 brand=quadrortx,driver>=535,driver<536 brand=nvidiartx,driver>=535,driver<536 brand=vapps,driver>=535,driver<536 brand=vpc,driver>=535,driver<536 brand=vcs,driver>=535,driver<536 brand=vws,driver>=535,driver<536 brand=cloudgaming,driver>=535,driver<536 brand=unknown,driver>=550,driver<551 brand=grid,driver>=550,driver<551 brand=tesla,driver>=550,driver<551 brand=nvidia,driver>=550,driver<551 brand=quadro,driver>=550,driver<551 brand=quadrortx,driver>=550,driver<551 brand=nvidiartx,driver>=550,driver<551 brand=vapps,driver>=550,driver<551 brand=vpc,driver>=550,driver<551 brand=vcs,driver>=550,driver<551 brand=vws,driver>=550,driver<551 brand=cloudgaming,driver>=550,driver<551
export NV_CUDA_CUDART_VERSION=12.6.77-1
export CUDA_VERSION=12.6.3
export LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
export NVIDIA_VISIBLE_DEVICES=all
export NVIDIA_DRIVER_CAPABILITIES=utility
export NVIDIA_DISABLE_REQUIRE=true
# 手动增加这个变量
export CONFIG_FILE=config/device-config.yaml
/root/go/bin/dlv debug --headless --listen=:12345 --api-version=2 ./cmd/device-plugin/nvidia/

3.3. 更改配置

由于节点的配置文件路径被固化在代码中，所以需要修改代码(todo: 这里可能可以优化)

diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go
--- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go
+++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go
@@ -101,7 +101,7 @@ type NvidiaDevicePlugin struct {
 }

 func readFromConfigFile(sConfig *nvidia.NvidiaConfig) (string, error) {
-       jsonbyte, err := os.ReadFile("/config/config.json")
+       jsonbyte, err := os.ReadFile("config/config.json")
        mode := "hami-core"
        if err != nil {
                return "", err
(END)

3.4. 日志等级调整

plugin的代码中日志都是以 klog的Warning, Error,Info为等级输出日志的, 所以无需关注日志等级，默认等级为0，所以三种等级的日志都会输出，无需进行额外的日志配置

plugin 启动命令中有个v，这个v不是日志等级，而是启动命令是否打印版本的配置

https://github.com/urfave/cli/blob/main/flag.go

// VersionFlag prints the version for the application
var VersionFlag Flag = &BoolFlag{
 Name:        "version",
 Aliases:     []string{"v"},
 Usage:       "print the version",
 HideDefault: true,
 Local:       true,
}

3.5. 关闭原来的服务

plugin以 daemonSet存在，所以要停止服务有以下方式可以参考

修改daemonSet镜像为一个不可达镜像，这样的话所有gpu节点的的plugin都会被停掉
修改要调试的节点label，与daemonSet的label不匹配即可

3.6. 启动新的服务

这里以远程方式连接到服务器上调试

在Goland中配置 Run/Debug Configurations 配置主机以及端口，与 start-plugin.sh 中的一只即可
在服务器上启动plugin服务

bash start-plugin.sh

这样就可以打断点调试plugin代码

4. hami-scheduler本地调试

4.1. 收集环境变量 hami-scheduler

 k exec hami-scheduler-67fc7ccd55-vjntl -c vgpu-scheduler-extender -- env > start-scheduler.sh

4.1.1. TLS配置

k8s admissionWebhook必须用https方式调用，所以启动的服务需要时https的。使用Helm Chart部署的hami服务已经创建好了证书，地址为127.0.0.1，所以启动的服务需要是 https://127.0.0.1:xxxx/webhook

需要收集TLS的CERT、KEY

k get secret hami-scheduler-tls -o jsonpath='{.data.tls\.crt}' |base64 -d > config/tls.crt
k get secret hami-scheduler-tls -o jsonpath='{.data.tls\.key}' |base64 -d > config/tls.key

编写启动命令, 并追加到 start-scheduler.sh

dlv debug --headless --listen=:2345 --api-version=2 ./cmd/scheduler/  -- \
  --device-config-file=config/device-config.yaml \
  -v=10 \
  --scheduler-name=hami-scheduler \
  --http_bind=0.0.0.0:8080 \
  --cert_file=config/tls.crt \
  --key_file=config/tls.key

最终得到一个启动脚本文件 start-scheduler.sh

注意PATH的处理，手动修改 PATH=$PATH:xxxxxxxx

PATH=$PATH:/k8s-vgpu/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
HOSTNAME=hami-scheduler-67fc7ccd55-vjntl
NVIDIA_MIG_MONITOR_DEVICES=all
HOOK_PATH=/usr/local
KUBERNETES_SERVICE_HOST=10.233.0.1
KUBERNETES_SERVICE_PORT=443
KUBERNETES_SERVICE_PORT_HTTPS=443
KUBERNETES_PORT=tcp://10.233.0.1:443
KUBERNETES_PORT_443_TCP=tcp://10.233.0.1:443
KUBERNETES_PORT_443_TCP_PROTO=tcp
KUBERNETES_PORT_443_TCP_PORT=443
KUBERNETES_PORT_443_TCP_ADDR=10.233.0.1
NVARCH=x86_64
NVIDIA_REQUIRE_CUDA=cuda>=12.6 brand=unknown,driver>=470,driver<471 brand=grid,driver>=470,driver<471 brand=tesla,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=quadro,driver>=470,driver<471 brand=quadrortx,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471 brand=vapps,driver>=470,driver<471 brand=vpc,driver>=470,driver<471 brand=vcs,driver>=470,driver<471 brand=vws,driver>=470,driver<471 brand=cloudgaming,driver>=470,driver<471 brand=unknown,driver>=535,driver<536 brand=grid,driver>=535,driver<536 brand=tesla,driver>=535,driver<536 brand=nvidia,driver>=535,driver<536 brand=quadro,driver>=535,driver<536 brand=quadrortx,driver>=535,driver<536 brand=nvidiartx,driver>=535,driver<536 brand=vapps,driver>=535,driver<536 brand=vpc,driver>=535,driver<536 brand=vcs,driver>=535,driver<536 brand=vws,driver>=535,driver<536 brand=cloudgaming,driver>=535,driver<536 brand=unknown,driver>=550,driver<551 brand=grid,driver>=550,driver<551 brand=tesla,driver>=550,driver<551 brand=nvidia,driver>=550,driver<551 brand=quadro,driver>=550,driver<551 brand=quadrortx,driver>=550,driver<551 brand=nvidiartx,driver>=550,driver<551 brand=vapps,driver>=550,driver<551 brand=vpc,driver>=550,driver<551 brand=vcs,driver>=550,driver<551 brand=vws,driver>=550,driver<551 brand=cloudgaming,driver>=550,driver<551
NV_CUDA_CUDART_VERSION=12.6.77-1
CUDA_VERSION=12.6.3
LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
NVIDIA_VISIBLE_DEVICES=all
NVIDIA_DRIVER_CAPABILITIES=utility
NVIDIA_DISABLE_REQUIRE=true

dlv debug --headless --listen=:2345 --api-version=2 ./cmd/scheduler/  -- \
  # device-config.yaml 使用之前从plugin中获取的即可
  --device-config-file=config/device-config.yaml \
  # 日志等级调高方便调试
  -v=10 \
  # 调度器名称需要
  --scheduler-name=hami-scheduler \
  # webhook需要使用的url
  --http_bind=0.0.0.0:8080 \
  --cert_file=config/tls.crt \
  --key_file=config/tls.key

4.2. webhook修改

如果要让pod启动时，kube-apiserver可以调用到,需要修改 admissionWebhook配置

k edit mutatingwebhookconfigurations.admissionregistration.k8s.io hami-webhook

原本的配置

webhooks:
- admissionReviewVersions:
  - v1beta1
  clientConfig:
    caBundle: LS0tLS1CRUdJTiBDRVJUSUZJQ0FURS0tLS0tCk1JSUJkakNDQVJ5Z0F3SUJBZ0lSQUxjd2FQMjUrMlphdGhTTlFMcG1qT0V3Q2dZSUtvWkl6ajBFQXdJd0R6RU4KTUFzR0ExVUVDaE1FYm1sc01UQWdGdzB5TkRFeU1EWXdOekV4TVRWYUdBOHlNVEkwTVRFeE1qQTNNVEV4TlZvdwpEekVOTUFzR0ExVUVDaE1FYm1sc01UQlpNQk1HQnlxR1NNNDlBZ0VHQ0NxR1NNNDlBd0VIQTBJQUJDUnlXUDdYCkRmT2N4NEVTMVRYaUs0dnFFU2wrcUFHYjI2YzNrOEdMWlZTL1lHaFpLZVVxaEgydVRhTFdWTW1hZVJFbkxqM0cKSStMVFRVTTR6SVhEUld5alZ6QlZNQTRHQTFVZER3RUIvd1FFQXdJQ0JEQVRCZ05WSFNVRUREQUtCZ2dyQmdFRgpCUWNEQVRBUEJnTlZIUk1CQWY4RUJUQURBUUgvTUIwR0ExVWREZ1FXQkJTcVV4bWpGa29YUlpRK0xXVzBNM1pJCnMzck1wakFLQmdncWhrak9QUVFEQWdOSUFEQkZBaUJSY2VRL2tJVkR2VTV3Vjl0K3NRWm93TmFhTWhIMTV5K2sKT3VrR0FlRGVtQUloQUxDZzFrM0JQZUJBNG8reWY5emxvVjM2VEk2RHUzaGdMT1B3MXhaZkFvcDMKLS0tLS1FTkQgQ0VSVElGSUNBVEUtLS0tLQo=
    service:
      name: hami-scheduler
      namespace: kube-system
      path: /webhook
      port: 443

修改为

webhooks:
- admissionReviewVersions:
  - v1beta1
  clientConfig:
    caBundle: LS0tLS1CRUdJTiBDRVJUSUZJQ0FURS0tLS0tCk1JSUJkakNDQVJ5Z0F3SUJBZ0lSQUxjd2FQMjUrMlphdGhTTlFMcG1qT0V3Q2dZSUtvWkl6ajBFQXdJd0R6RU4KTUFzR0ExVUVDaE1FYm1sc01UQWdGdzB5TkRFeU1EWXdOekV4TVRWYUdBOHlNVEkwTVRFeE1qQTNNVEV4TlZvdwpEekVOTUFzR0ExVUVDaE1FYm1sc01UQlpNQk1HQnlxR1NNNDlBZ0VHQ0NxR1NNNDlBd0VIQTBJQUJDUnlXUDdYCkRmT2N4NEVTMVRYaUs0dnFFU2wrcUFHYjI2YzNrOEdMWlZTL1lHaFpLZVVxaEgydVRhTFdWTW1hZVJFbkxqM0cKSStMVFRVTTR6SVhEUld5alZ6QlZNQTRHQTFVZER3RUIvd1FFQXdJQ0JEQVRCZ05WSFNVRUREQUtCZ2dyQmdFRgpCUWNEQVRBUEJnTlZIUk1CQWY4RUJUQURBUUgvTUIwR0ExVWREZ1FXQkJTcVV4bWpGa29YUlpRK0xXVzBNM1pJCnMzck1wakFLQmdncWhrak9QUVFEQWdOSUFEQkZBaUJSY2VRL2tJVkR2VTV3Vjl0K3NRWm93TmFhTWhIMTV5K2sKT3VrR0FlRGVtQUloQUxDZzFrM0JQZUJBNG8reWY5emxvVjM2VEk2RHUzaGdMT1B3MXhaZkFvcDMKLS0tLS1FTkQgQ0VSVElGSUNBVEUtLS0tLQo=
    url: https://127.0.0.1:8080/webhook

4.3. 为拓展的hami-scheduler启动一个HTTP服务

由于上面的TLS相关配置都是127.0.0.1的，admissionWebhook是从kube-apiserver POD中调用的，使用的本机ip，所以可以通信，但是hami-scheduler启动了一个新的schedler在POD中使用了虚拟网络所以无法掉通以进程方式直接启动在HOST上的，所以我们需要为hami-scheduler启动一个新的可以掉通的HTTP服务，这里有多种方式，例如转发，或者多启动一个HTTP SERVER，这里选择第二种方式, 需要修改调度器代码

diff --git a/cmd/scheduler/main.go b/cmd/scheduler/main.go
index 3c99042..c85c09f 100644
--- a/cmd/scheduler/main.go
+++ b/cmd/scheduler/main.go
@@ -84,12 +85,29 @@ func start() {
        router.GET("/healthz", routes.HealthzRoute())
        klog.Info("listen on ", config.HTTPBind)
        if len(tlsCertFile) == 0 || len(tlsKeyFile) == 0 {
-               if err := http.ListenAndServe(config.HTTPBind, router); err != nil {
-                       klog.Fatal("Listen and Serve error, ", err)
+               go func() {
+                       if err := http.ListenAndServe(config.HTTPBind, router); err != nil {
+                               klog.Fatal("Listen and Serve error, ", err)
+                       }
+               }()
+       } else {
+               go func() {
+                       if err := http.ListenAndServeTLS(config.HTTPBind, tlsCertFile, tlsKeyFile, router); err != nil {
+                               klog.Fatal("Listen and Serve TLS error, ", err)
+                       }
+               }()
+       }
+
+       // Additional HTTP server on a different port
+       additionalBind := "10.10.10.8:8081" 
+       klog.Info("listen on ", additionalBind)
+       if len(tlsCertFile) == 0 || len(tlsKeyFile) == 0 {
+               if err := http.ListenAndServe(additionalBind, router); err != nil {
+                       klog.Fatal("Additional Listen and Serve error, ", err)
                }
        } else {
-               if err := http.ListenAndServeTLS(config.HTTPBind, tlsCertFile, tlsKeyFile, router); err != nil {
-                       klog.Fatal("Listen and Serve error, ", err)
+               if err := http.ListenAndServeTLS(additionalBind, tlsCertFile, tlsKeyFile, router); err != nil {
+                       klog.Fatal("Additional Listen and Serve TLS error, ", err)
                }
        }
 }

这样就启动了两个不同HOST+PORT的服务, adminission webhook用 https://127.0.0.1:xxxx 拓展调度器用 https://10.10.10.8:8081/xxx

4.3.1. 修改scheduler服务器地址

拓展调度器的服务端准备好了，客户端需要做如下修改

k edit cm hami-scheduler-newversion

修改

data:
  config.yaml: |
  ...
    extenders:
    - urlPrefix: "https://127.0.0.1:8081"

将 urlPrefix 修改为刚刚启动的服务路径

data:
  config.yaml: |
  ...
    extenders:
    - urlPrefix: "https://10.10.10.8:8081"

4.4. 关闭原来的hami-scheduler

deployment hami-scheduler 中有两个容器，

kube-scheduler 拓展原有的k8s调度器，将 schedulerName 为 hami-scheduler 的pod的调度处理转发到 vgpu-scheduler-extender 中
vgpu-scheduler-extender 具体的调度器实现逻辑

其中 kube-scheduler 容器只是作为拓展调度器的注册以及请求转发，需要调试的地方不多，所以使用当前存在的即可，vgpu-scheduler-extender 需要停止，所以最简单的方式还是修改这个容器镜像地址为不可达即可

4.5. 启动本地的scheduler

bash start-scheduler.sh

在Goland中做类似 hami-plugin 配置，即可远程调试

注意dlv启动的端口需要和启动脚本中的保持一致，如果同事需要多个远程调试进程则建议将端口改为不同的

5. 测试

创建一个pod，看调度流程是否正常，如果可以调度成功，且日志正常即可，测试通过后即可进行功能开发或者bug修复

k apply -f examples/nvidia/default_use.yaml

HAMI Project Local Debugging

Fri, 20 Dec 2024 00:00:00 +0000

English version generated by AI. 查看中文版

HAMI Project Local Debugging

HAMI Project Local Debugging

1. Prerequisites

A fully installed and functional HAMI cluster on a k8s cluster is required for debugging, allowing for the reuse of many configurations.

HAMI version: v2.4.1
OS: Ubuntu 22.04.4 LTS
Kernel: 5.15.0-125-generic
NVCC: V11.5.119
Nvidia-smi: 550.54.14
CUDA: 12.4

nvidia-smi
Fri Dec 20 02:45:30 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.14              Driver Version: 550.54.14      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3090        Off |   00000000:00:06.0 Off |                  N/A |
|  0%   26C    P8             11W /  370W |       0MiB /  24576MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

2. Create Local Directory for Configuration

Local startup depends on some files, which can be stored in a unified folder.

2.1. Create Folder

For example, create a config directory under the project.

mkdir config

3. Local Debugging of hami-device

3.1. Collect Configuration

hami-device requires two configuration files:

config/device-config.yaml
config.json

Retrieve and save them to the local directory.

k exec hami-device-plugin-fn8lc -c device-plugin -- cat /device-config.yaml > /config/device.yaml
k exec hami-device-plugin-fn8lc -c device-plugin -- cat /config/config.json > /config/config.json

3.2. Collect Environment Variables

tip: before exporting pod environment variables, configure enableServiceLinks for the pod, such as setting spec.template.spec.enableServiceLinks: false in the deployment configuration. This allows you to remove the information automatically added by k8s and only keep the environment variables configured by the pod itself.

k exec hami-device-plugin-lzm4n env > start-plugin.sh

Note: The exported file contains a PATH variable, which needs to be changed to export PATH=$PATH:xxxxxx with xxx replaced by the original output; add export in front of each line. Use dlv with goland for debugging.

Remove configurations containing IP addresses; retain port configurations; add debug information to obtain a file as follows:

#!/usr/bin/bash
export PATH=$PATH/k8s-vgpu/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
export HOSTNAME=hami-scheduler-858f6b9dcf-4nn58
export KUBERNETES_SERVICE_PORT=443
export KUBERNETES_SERVICE_PORT_HTTPS=443
export KUBERNETES_PORT=tcp://10.233.0.1:443
export KUBERNETES_PORT_443_TCP=tcp://10.233.0.1:443
export KUBERNETES_PORT_443_TCP_PROTO=tcp
export KUBERNETES_PORT_443_TCP_PORT=443
export KUBERNETES_PORT_443_TCP_ADDR=10.233.0.1
export KUBERNETES_SERVICE_HOST=10.233.0.1
export NVARCH=x86_64
export NVIDIA_REQUIRE_CUDA=cuda>=12.6 brand=unknown,driver>=470,driver<471 brand=grid,driver>=470,driver<471 brand=tesla,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=quadro,driver>=470,driver<471 brand=quadrortx,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471 brand=vapps,driver>=470,driver<471 brand=vpc,driver>=470,driver<471 brand=vcs,driver>=470,driver<471 brand=vws,driver>=470,driver<471 brand=cloudgaming,driver>=470,driver<471 brand=unknown,driver>=535,driver<536 brand=grid,driver>=535,driver<536 brand=tesla,driver>=535,driver<536 brand=nvidia,driver>=535,driver<536 brand=quadro,driver>=535,driver<536 brand=quadrortx,driver>=535,driver<536 brand=nvidiartx,driver>=535,driver<536 brand=vapps,driver>=535,driver<536 brand=vpc,driver>=535,driver<536 brand=vcs,driver>=535,driver<536 brand=vws,driver>=535,driver<536 brand=cloudgaming,driver>=535,driver<536 brand=unknown,driver>=550,driver<551 brand=grid,driver>=550,driver<551 brand=tesla,driver>=550,driver<551 brand=nvidia,driver>=550,driver<551 brand=quadro,driver>=550,driver<551 brand=quadrortx,driver>=550,driver<551 brand=nvidiartx,driver>=550,driver<551 brand=vapps,driver>=550,driver<551 brand=vpc,driver>=550,driver<551 brand=vcs,driver>=550,driver<551 brand=vws,driver>=550,driver<551 brand=cloudgaming,driver>=550,driver<551
export NV_CUDA_CUDART_VERSION=12.6.77-1
export CUDA_VERSION=12.6.3
export LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
export NVIDIA_VISIBLE_DEVICES=all
export NVIDIA_DRIVER_CAPABILITIES=utility
export NVIDIA_DISABLE_REQUIRE=true
# Manually add this variable
export CONFIG_FILE=config/device-config.yaml
/root/go/bin/dlv debug --headless --listen=:12345 --api-version=2 ./cmd/device-plugin/nvidia/

3.3. Modify Configuration

Since the node’s configuration file path is hardcoded in the code, it needs to be modified (todo: this could potentially be optimized).

diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go
--- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go
+++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go
@@ -101,7 +101,7 @@ type NvidiaDevicePlugin struct {
 }

 func readFromConfigFile(sConfig *nvidia.NvidiaConfig) (string, error) {
-       jsonbyte, err := os.ReadFile("/config/config.json")
+       jsonbyte, err := os.ReadFile("config/config.json")
        mode := "hami-core"
        if err != nil {
                return "", err
(END)

3.4. Adjust Log Level

The plugin’s code outputs logs at the levels of klog’s Warning, Error, Info, so there is no need to worry about log levels. The default level is 0, so all three levels of logs will be output without additional log configuration.

The v in the plugin startup command is not a log level, but a configuration for whether to print the version.

https://github.com/urfave/cli/blob/main/flag.go

// VersionFlag prints the version for the application
var VersionFlag Flag = &BoolFlag{
 Name:        "version",
 Aliases:     []string{"v"},
 Usage:       "print the version",
 HideDefault: true,
 Local:       true,
}

3.5. Stop the Original Service

The plugin exists as a daemonSet, so there are several ways to stop the service:

Modify the daemonSet image to an unreachable image, which will stop all GPU node plugins.
Modify the label of the node to be debugged so that it does not match the daemonSet’s label.

3.6. Start the New Service

Connect to the server remotely for debugging:

Configure the host and port in Goland’s Run/Debug Configurations to match those in start-plugin.sh.
Start the plugin service on the server.

bash start-plugin.sh

This allows for breakpoint debugging of the plugin code.

4. Local Debugging of hami-scheduler

4.1. Collect Environment Variables for hami-scheduler

 k exec hami-scheduler-67fc7ccd55-vjntl -c vgpu-scheduler-extender -- env > start-scheduler.sh

4.1.1. TLS Configuration

The k8s admissionWebhook must be called using HTTPS, so the started service needs to be HTTPS. The Helm Chart-deployed hami service has already created the certificate, with the address at 127.0.0.1, so the started service needs to be https://127.0.0.1:xxxx/webhook.

Collect the TLS CERT and KEY:

k get secret hami-scheduler-tls -o jsonpath='{.data.tls\.crt}' |base64 -d > config/tls.crt
k get secret hami-scheduler-tls -o jsonpath='{.data.tls\.key}' |base64 -d > config/tls.key

Write the startup command and append it to start-scheduler.sh:

dlv debug --headless --listen=:2345 --api-version=2 ./cmd/scheduler/  -- \
  --device-config-file=config/device-config.yaml \
  -v=10 \
  --scheduler-name=hami-scheduler \
  --http_bind=0.0.0.0:8080 \
  --cert_file=config/tls.crt \
  --key_file=config/tls.key

This results in a startup script file start-scheduler.sh.

Note the handling of the PATH, manually modify PATH=$PATH:xxxxxxxx

PATH=$PATH:/k8s-vgpu/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
HOSTNAME=hami-scheduler-67fc7ccd55-vjntl
NVIDIA_MIG_MONITOR_DEVICES=all
HOOK_PATH=/usr/local
KUBERNETES_SERVICE_HOST=10.233.0.1
KUBERNETES_SERVICE_PORT=443
KUBERNETES_SERVICE_PORT_HTTPS=443
KUBERNETES_PORT=tcp://10.233.0.1:443
KUBERNETES_PORT_443_TCP=tcp://10.233.0.1:443
KUBERNETES_PORT_443_TCP_PROTO=tcp
KUBERNETES_PORT_443_TCP_PORT=443
KUBERNETES_PORT_443_TCP_ADDR=10.233.0.1
NVARCH=x86_64
NVIDIA_REQUIRE_CUDA=cuda>=12.6 brand=unknown,driver>=470,driver<471 brand=grid,driver>=470,driver<471 brand=tesla,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=quadro,driver>=470,driver<471 brand=quadrortx,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471 brand=vapps,driver>=470,driver<471 brand=vpc,driver>=470,driver<471 brand=vcs,driver>=470,driver<471 brand=vws,driver>=470,driver<471 brand=cloudgaming,driver>=470,driver<471 brand=unknown,driver>=535,driver<536 brand=grid,driver>=535,driver<536 brand=tesla,driver>=535,driver<536 brand=nvidia,driver>=535,driver<536 brand=quadro,driver>=535,driver<536 brand=quadrortx,driver>=535,driver<536 brand=nvidiartx,driver>=535,driver<536 brand=vapps,driver>=535,driver<536 brand=vpc,driver>=535,driver<536 brand=vcs,driver>=535,driver<536 brand=vws,driver>=535,driver<536 brand=cloudgaming,driver>=535,driver<536 brand=unknown,driver>=550,driver<551 brand=grid,driver>=550,driver<551 brand=tesla,driver>=550,driver<551 brand=nvidia,driver>=550,driver<551 brand=quadro,driver>=550,driver<551 brand=quadrortx,driver>=550,driver<551 brand=nvidiartx,driver>=550,driver<551 brand=vapps,driver>=550,driver<551 brand=vpc,driver>=550,driver<551 brand=vcs,driver>=550,driver<551 brand=vws,driver>=550,driver<551 brand=cloudgaming,driver>=550,driver<551
NV_CUDA_CUDART_VERSION=12.6.77-1
CUDA_VERSION=12.6.3
LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
NVIDIA_VISIBLE_DEVICES=all
NVIDIA_DRIVER_CAPABILITIES=utility
NVIDIA_DISABLE_REQUIRE=true

dlv debug --headless --listen=:2345 --api-version=2 ./cmd/scheduler/  -- \
 # Use the previously obtained device-config.yaml
  --device-config-file=config/device-config.yaml \
  # Increase log level for easier debugging
  -v=10 \
  # Scheduler name is required
  --scheduler-name=hami-scheduler \
  # URL needed for the webhook
  --http_bind=0.0.0.0:8080 \
  --cert_file=config/tls.crt \
  --key_file=config/tls.key

4.2. Modify Webhook

To allow the kube-apiserver to call the pod at startup, modify the admissionWebhook configuration.

k edit mutatingwebhookconfigurations.admissionregistration.k8s.io hami-webhook

Original configuration:

webhooks:
- admissionReviewVersions:
  - v1beta1
  clientConfig:
    caBundle: LS0tLS1CRUdJTiBDRVJUSUZJQ0FURS0tLS0tCk1JSUJkakNDQVJ5Z0F3SUJBZ0lSQUxjd2FQMjUrMlphdGhTTlFMcG1qT0V3Q2dZSUtvWkl6ajBFQXdJd0R6RU4KTUFzR0ExVUVDaE1FYm1sc01UQWdGdzB5TkRFeU1EWXdOekV4TVRWYUdBOHlNVEkwTVRFeE1qQTNNVEV4TlZvdwpEekVOTUFzR0ExVUVDaE1FYm1sc01UQlpNQk1HQnlxR1NNNDlBZ0VHQ0NxR1NNNDlBd0VIQTBJQUJDUnlXUDdYCkRmT2N4NEVTMVRYaUs0dnFFU2wrcUFHYjI2YzNrOEdMWlZTL1lHaFpLZVVxaEgydVRhTFdWTW1hZVJFbkxqM0cKSStMVFRVTTR6SVhEUld5alZ6QlZNQTRHQTFVZER3RUIvd1FFQXdJQ0JEQVRCZ05WSFNVRUREQUtCZ2dyQmdFRgpCUWNEQVRBUEJnTlZIUk1CQWY4RUJUQURBUUgvTUIwR0ExVWREZ1FXQkJTcVV4bWpGa29YUlpRK0xXVzBNM1pJCnMzck1wakFLQmdncWhrak9QUVFEQWdOSUFEQkZBaUJSY2VRL2tJVkR2VTV3Vjl0K3NRWm93TmFhTWhIMTV5K2sKT3VrR0FlRGVtQUloQUxDZzFrM0JQZUJBNG8reWY5emxvVjM2VEk2RHUzaGdMT1B3MXhaZkFvcDMKLS0tLS1FTkQgQ0VSVElGSUNBVEUtLS0tLQo=
    service:
      name: hami-scheduler
      namespace: kube-system
      path: /webhook
      port: 443

Modify the configuration to point to the new service address:

webhooks:
- admissionReviewVersions:
  - v1beta1
  clientConfig:
    url: https://127.0.0.1:8080/webhook
    caBundle: LS0tLS1CRUdJTiBDRVJUSUZJQ0FURS0tLS0tCk1JSUJkakNDQVJ5Z0F3SUJBZ0lSQUxjd2FQMjUrMlphdGhTTlFMcG1qT0V3Q2dZSUtvWkl6ajBFQXdJd0R6RU4KTUFzR0ExVUVDaE1FYm1sc01UQWdGdzB5TkRFeU1EWXdOekV4TVRWYUdBOHlNVEkwTVRFeE1qQTNNVEV4TlZvdwpEekVOTUFzR0ExVUVDaE1FYm1sc01UQlpNQk1HQnlxR1NNNDlBZ0VHQ0NxR1NNNDlBd0VIQTBJQUJDUnlXUDdYCkRmT2N4NEVTMVRYaUs0dnFFU2wrcUFHYjI2YzNrOEdMWlZTL1lHaFpLZVVxaEgydVRhTFdWTW1hZVJFbkxqM0cKSStMVFRVTTR6SVhEUld5alZ6QlZNQTRHQTFVZER3RUIvd1FFQXdJQ0JEQVRCZ05WSFNVRUREQUtCZ2dyQmdFRgpCUWNEQVRBUEJnTlZIUk1CQWY4RUJUQURBUUgvTUIwR0ExVWREZ1FXQkJTcVV4bWpGa29YUlpRK0xXVzBNM1pJCnMzck1wakFLQmdncWhrak9QUVFEQWdOSUFEQkZBaUJSY2VRL2tJVkR2VTV3Vjl0K3NRWm93TmFhTWhIMTV5K2sKT3VrR0FlRGVtQUloQUxDZzFrM0JQZUJBNG8reWY5emxvVjM2VEk2RHUzaGdMT1B3MXhaZkFvcDMKLS0tLS1FTkQgQ0VSVElGSUNBVEUtLS0tLQo=
    service:
      name: hami-scheduler
      namespace: kube-system
      path: /webhook
      port: 443

4.3. Start an HTTP Service for the Extended hami-scheduler

Since the TLS configuration is set to only accept connections from 127.0.0.1, it’s necessary to start an HTTP service that can be accessed from the kube-apiserver. This involves modifying the scheduler code to listen on an additional port.

diff --git a/cmd/scheduler/main.go b/cmd/scheduler/main.go
index 3c99042..c85c09f 100644
--- a/cmd/scheduler/main.go
+++ b/cmd/scheduler/main.go
@@ -84,12 +85,29 @@ func start() {
        router.GET("/healthz", routes.HealthzRoute())
        klog.Info("listen on ", config.HTTPBind)
        if len(tlsCertFile) == 0 || len(tlsKeyFile) == 0 {
-               if err := http.ListenAndServe(config.HTTPBind, router); err != nil {
-                       klog.Fatal("Listen and Serve error, ", err)
+               go func() {
+                       if err := http.ListenAndServe(config.HTTPBind, router); err != nil {
+                               klog.Fatal("Listen and Serve error, ", err)
+                       }
+               }()
+       } else {
+               go func() {
+                       if err := http.ListenAndServeTLS(config.HTTPBind, tlsCertFile, tlsKeyFile, router); err != nil {
+                               klog.Fatal("Listen and Serve TLS error, ", err)
+                       }
+               }()
+       }
+
+       // Additional HTTP server on a different port
+       additionalBind := "0.0.0.0:8081" 
+       klog.Info("listen on ", additionalBind)
+       if len(tlsCertFile) == 0 || len(tlsKeyFile) == 0 {
+               if err := http.ListenAndServe(additionalBind, router); err != nil {
+                       klog.Fatal("Additional Listen and Serve error, ", err)
                }
        } else {
-               if err := http.ListenAndServeTLS(config.HTTPBind, tlsCertFile, tlsKeyFile, router); err != nil {
-                       klog.Fatal("Listen and Serve error, ", err)
+               if err := http.ListenAndServeTLS(additionalBind, tlsCertFile, tlsKeyFile, router); err != nil {
+                       klog.Fatal("Additional Listen and Serve TLS error, ", err)
                }
        }
 }

This modification starts an additional HTTP server that listens on a different port, allowing the kube-apiserver to communicate with the scheduler via a secure connection.

4.4. Stop the Original hami-scheduler

To debug the new scheduler setup, the original running instance of the hami-scheduler needs to be stopped. This can be done by scaling down the deployment or modifying the deployment to use an image that does not start the scheduler.

kubectl scale deployment hami-scheduler --replicas=0

4.5. Start the Local Scheduler

With the environment prepared, start the local version of the scheduler using the script prepared earlier.

bash start-scheduler.sh

This command starts the scheduler in debug mode, allowing for real-time debugging and testing of changes.

5. Testing

To verify that the local debugging setup works correctly, deploy a test pod that requires scheduling by the hami-scheduler.

kubectl apply -f test-pod.yaml

Monitor the scheduler’s output and the Kubernetes events to ensure the pod is scheduled correctly and that the scheduler behaves as expected.

By following these steps, local debugging of the HAMI project can be effectively carried out, allowing for rapid development and troubleshooting of the scheduling components.