深度解析 DeepFlow 如何采集大模型服务的业务指标

type llmParser struct {
    httpStream map[uint64] *StreamInfo
}
type StreamInfo struct {
    reqTime uint64    // 请求时间
    respFirstChunkedTime uint64   //首次分块响应的时间
    totalToken uint64  // 总的输出Token数量
    flag int  //标记是否已经记录了首次分块响应的时间
}

func checker(payload []byte) (protoNum uint8, protoStr string) {
    // 读取HTTP请求。
    req, err := http.ReadRequest(bufio.NewReader(bytes.NewReader(payload)))
    if err != nil {
        return 0, ""
    }
    query := req.URL.Path
    // 检查请求路径是否包含/generate_stream，
    if strings.Contains(query, "/generate_stream") {
        sdk.Warn(fmt.Sprintf("check: %s", query))
        // 如果包含，返回协议编号1和字符串http_stream，
        return 1, "http_stream"
    }
    return 0, "" // 否则返回0和空字符串。
}

func (p *llmParser) HookIn() []sdk.HookBitmap {
    return []sdk.HookBitmap{
        sdk.HOOK_POINT_PAYLOAD_PARSE,
    }
}

func (p *llmParser) OnCheckPayload(baseCtx *sdk.ParseCtx) (protoNum uint8, protoStr string) {
    if baseCtx.EbpfType != sdk.EbpfTypeNone {
        return 0, ""
    }
    // 获取负载数据
    payload, err := baseCtx.GetPayload()
    if err != nil {
        return 0, ""
    }
    // 如果是请求方向，调用 checker 函数进行检查
    if baseCtx.Direction == sdk.DirectionRequest {
        return checker(payload)
    }
    return 0, ""
}

func (p *llmParser) OnParsePayload(baseCtx *sdk.ParseCtx) sdk.Action {
    // 检查 baseCtx.L7 是否为 1。
    if baseCtx.L7 != 1 {
        return sdk.ActionNext()
    }
    // 尝试从 baseCtx 中获取 payload
    payload, err := baseCtx.GetPayload()
    if err != nil {
        return sdk.ActionAbortWithErr(err)
    }
    var attr = []sdk.KeyVal{}
    // 获取当前流的 ID flowId。
    var flowId = baseCtx.FlowID
    // 检查 p.httpStream 中是否存在该流的信息，如果不存在则初始化一个新的 StreamInfo 对象。
    if p.httpStream[flowId] == nil {
        p.httpStream[flowId] = &StreamInfo{}
    }
    switch baseCtx.Direction {
    case sdk.DirectionRequest: // 如果 baseCtx.Direction 是 sdk.DirectionRequest，则处理
        req, err := http.ReadRequest(bufio.NewReader(bytes.NewReader(payload)))
        if err != nil {
            return sdk.ActionNext()
        }
        p.httpStream[flowId].reqTime = baseCtx.Time // 记录请求时间 reqTime
        info := &sdk.L7ProtocolInfo{
            Req: &sdk.Request{
                Resource: req.URL.Path,
            },
            Resp: &sdk.Response{},
        }
        // 返回 sdk.ParseActionAbortWithL7Info，表示终止处理并返回 L7 协议信息
        return sdk.ParseActionAbortWithL7Info([]*sdk.L7ProtocolInfo{info})
    case sdk.DirectionResponse:
        // 使用 bufio.Reader 读取响应数据
        r := bufio.NewReader(bytes.NewReader(payload))
        bs, _, err := r.ReadLine()
        if err == io.EOF {
            return sdk.ActionNext()
        }
        regex := regexp.MustCompile(`^HTTP/[1-2]\.[01] \d{3} .*$`)
        if regex.MatchString(string(bs)) {
            return sdk.ActionNext()
        }
        // 如果是结束标志（0），计算并记录时间差和每个 token 的时间
        if string(bs) == "0" {
            attr = []sdk.KeyVal{
                {
                    // 首 Token 响应时间 - 请求发出时间
                    Key: "ttft",
                    Val: fmt.Sprintf("%d", p.httpStream[flowId].respFirstChunkedTime-p.httpStream[flowId].reqTime),
                },
                {
                    // （结束时间 - 首 Token 响应时间）/ 总 Token 数
                    Key: "tpot",
                    Val: fmt.Sprintf("%d", (baseCtx.Time-p.httpStream[flowId].respFirstChunkedTime)/p.httpStream[flowId].totalToken),
                },
            }
            info := &sdk.L7ProtocolInfo{
                Req:  &sdk.Request{},
                Resp: &sdk.Response{},
                Kv:   attr,
            }
            if _, exists := p.httpStream[flowId]; exists {
                delete(p.httpStream, flowId)
            }
            return sdk.ParseActionAbortWithL7Info([]*sdk.L7ProtocolInfo{info})
        }
        bs, _, err = r.ReadLine()
        if err == io.EOF {
            return sdk.ActionNext()
        }
        // 如果是响应首包，记录首次分块时间和 token 数
        if p.httpStream[flowId].flag == 0 {
            p.httpStream[flowId].flag = 1
            p.httpStream[flowId].respFirstChunkedTime = baseCtx.Time
            p.httpStream[flowId].totalToken = uint64(len(bs))
            return sdk.ActionNext()
        }
        // 继续读取后续分块，累加 token 数
        p.httpStream[flowId].totalToken += uint64(len(bs))
        bs, _, err = r.ReadLine()
        // 如果读取到 EOF，删除流信息并返回
        if err == io.EOF {
            if _, exists := p.httpStream[flowId]; exists {
                delete(p.httpStream, flowId)
            }
            return sdk.ActionNext()
        }
        return sdk.ActionNext()
    default:
        return sdk.ActionNext()
    }
}

func main() {
    sdk.Warn("llm wasm plugin loaded") // 输出日志信息，表示插件已加载。
    llm := &llmParser{ // 创建 llmParser 实例
        httpStream: map[uint64]*StreamInfo{},
    }
    sdk.SetParser(llm) // 设置解析器为 llmParser 实例
}

编译：tinygo build -o llm.wasm -target wasi -gc=precise -panic=trap -scheduler=none -no-debug
上传：deepflow-ctl plugin create --type wasm --image llm.wasm --name llm ./llm/llm.go
查看：deepflow-ctl plugin list