突发流量处理最佳实践

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21import os
from arkruntime import Ark

client = Ark(
    api_key=os.environ.get("ARK_API_KEY"),
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 非流式输出：超时时间 = 排队时间 + 基础超时时间 = 300 秒 + 1800秒
    timeout=2100,
)

completion = client.chat.completions.create(
    # Replace with Model ID
    model="doubao-seed-2-1-pro-260915",
    messages = [
        {"role": "user", "content": "常见的十字花科植物有哪些？"},
    ],
    # 自定义request，突发流量最大排队时间 300000 毫秒（5分钟）
    extra_headers={"X-Ark-Max-Wait-Timeout-Ms": "300000"},
    stream = False,
)
print(completion.choices[0].message.content)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74package main

import (
  "context"
  "fmt"
  "log"
  "os"
  "time"

  "github.com/volcengine/ark-runtime-go/arkruntime"
  "github.com/volcengine/ark-runtime-go/arkruntime/model/chat"
)

const (
  ArkModel   = "doubao-seed-2-1-pro-260915"
  ArkBaseUrl = "https://ark.cn-beijing.volces.com/api/v3"
)

var (
  ArkApiKey = os.Getenv("ARK_API_KEY")
)

func callArk(msg string) (string, error) {
  client := arkruntime.NewClientWithApiKey(
    ArkApiKey,
    arkruntime.WithBaseUrl(ArkBaseUrl),
  )

  // 非流式输出：超时时间 = 排队时间 + 基础超时时间 = 300 秒 + 1800秒
  ctx, cancel := context.WithTimeout(context.Background(), 2100*time.Second)
  defer cancel()

  resp, err := client.CreateChatCompletion(
    ctx,
    &chat.ChatCompletionRequest{
      Model: ArkModel,
      Messages: []chat.ChatCompletionRequestMessage{
userMessage(msg),
      },
    },
    arkruntime.WithCustomHeaders(map[string]string{
      // 自定义request，突发流量最大排队时间 300000 毫秒 （5分钟）
      "X-Ark-Max-Wait-Timeout-Ms": "300000",
    }),
  )
  if err != nil {
    return "", err
  }

  if len(resp.Choices) > 0 && resp.Choices[0].Message.Content != "" {
    return resp.Choices[0].Message.Content, nil
  }
  return "", fmt.Errorf("empty response")
}

func main() {
  result, err := callArk("常见的十字花科植物有哪些？")
  if err != nil {
    log.Fatalf("调用失败: %v", err)
  }

  fmt.Println(result)
}

func userMessage(content string) chat.ChatCompletionRequestMessage {
    return chat.ChatCompletionRequestMessage{
        OneOf: chat.NewChatCompletionRequestUserMessageChatCompletionRequestMessageSum(
            chat.ChatCompletionRequestUserMessage{
                Role: chat.ChatCompletionRequestUserMessageRoleUser,
                Content: chat.NewOptNilChatCompletionMessageContent(chat.NewStringChatCompletionMessageContent(content)),
            },
        ),
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61package com.ark.sample;

import com.volcengine.ark.runtime.models.chat.*;
import com.volcengine.ark.runtime.service.ArkService;

import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Header {
    private static final String ARK_MODEL = "doubao-seed-2-1-pro-260915";
    private static final String ARK_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3";
    private static final String ARK_API_KEY = System.getenv("ARK_API_KEY");

    public static String callArk(String msg) {
        ArkService service = ArkService.builder()
                .apiKey(ARK_API_KEY)
                .baseUrl(ARK_BASE_URL)
                // 非流式输出：超时时间 = 排队时间 + 基础超时时间 = 300 秒 + 1800秒
                .timeout(Duration.ofSeconds(2100))
                .build();

        final List<ChatCompletionRequestMessage> messages = new ArrayList<>();
        final ChatCompletionRequestMessage userMessage = ChatCompletionRequestUserMessage.builder()
                .role(ChatCompletionRequestMessageType.USER)
                .content(ChatCompletionMessageContent.ofString(msg))
                .build();
        messages.add(userMessage);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(ARK_MODEL)
                .messages(messages)
                .build();

        Map<String, String> customHeaders = new HashMap<>();
        // 自定义request，突发流量最大排队时间 300000 毫秒 （5分钟）
        customHeaders.put("X-Ark-Max-Wait-Timeout-Ms", "300000");

        String result = (String) service.createChatCompletion(chatCompletionRequest, customHeaders)
                .getChoices()
                .get(0)
                .getMessage()
                .getContent();


        service.shutdownExecutor();
        return result;
    }

    public static void main(String[] args) {
        try {
            String result = callArk("常见的十字花科植物有哪些？");
            System.out.println(result);
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import time
import json
import os
import threading
# 安装SDK:  python -m pip install --upgrade arkruntime pika
from arkruntime import Ark
import pika

# ===================== 全局极简配置【所有参数可一键修改】=====================
MQ_QUEUE = "client_ark_queue"  # MQ削峰队列
MQ_CONF = pika.ConnectionParameters('127.0.0.1', 5672, credentials=pika.PlainCredentials('guest', 'guest'))
ARK_API_KEY = os.getenv('ARK_API_KEY')  # 获取API Key: https://ark.volcengine.com/region:cn-beijing/apikey
ARK_MODEL = "doubao-seed-2-1-pro-260915"
# 🔥 斜率可控核心配置 🔥 斜率=步长/间隔 = 2/5 = 0.4QPS/秒，绝对无陡增
FLOW = {"curr_qps": 5, "max_qps": 30, "step": 2, "interval": 5, "count": 0}
LOCK = threading.Lock()

# ===================== 1. 极简斜率可控流控【核心】=====================
def flow_schedule():
    """后台线程：匀速增长消费速率，控制斜率"""
    while True:
        time.sleep(FLOW["interval"])
        if FLOW["curr_qps"] < FLOW["max_qps"]:
            FLOW["curr_qps"] += FLOW["step"]

def reset_count():
    """后台线程：每秒重置请求计数器，精准控QPS"""
    while True:
        time.sleep(1)
        with LOCK: FLOW["count"] = 0

def allow():
    """流控校验：是否允许发起本次请求"""
    with LOCK:
        if FLOW["count"] < FLOW["curr_qps"]:
            FLOW["count"] += 1
            return True
    return False

# ===================== 2. MQ生产者【Client端请求投递，削峰】=====================
def producer(client_req):
    """投递Client请求到MQ，无速率限制，承接峰值流量"""
    conn = pika.BlockingConnection(MQ_CONF)
    ch = conn.channel()
    ch.queue_declare(queue=MQ_QUEUE, durable=True)  # 队列持久化
    ch.basic_publish(exchange='', routing_key=MQ_QUEUE, body=json.dumps(client_req),
                     properties=pika.BasicProperties(delivery_mode=2))  # 消息持久化
    conn.close()

# ===================== 3. 方舟SDK调用【大模型请求核心】=====================
def call_ark(msg):
    """✅ 标准使用指定方舟SDK，极简调用大模型"""
    client = Ark(
        api_key=ARK_API_KEY,
        base_url='https://ark.cn-beijing.volces.com/api/v3',
    )
    resp = client.chat.completions.create(model=ARK_MODEL, messages=[{"role": "user", "content": msg}])
    return resp.choices[0].message.content

# ===================== 4. MQ消费者【填谷+斜率控速+调用方舟，核心主逻辑】=====================
def consumer():
    """消费端核心：削峰后的流量，斜率可控匀速消费，调用方舟大模型"""
    conn = pika.BlockingConnection(MQ_CONF)
    ch = conn.channel()
    ch.queue_declare(queue=MQ_QUEUE, durable=True)
    ch.basic_qos(prefetch_count=1)  # 防止消费积压

    def callback(_ch, method, _, body):
        req = json.loads(body)["msg"]
        if not allow():  # 流控校验，不通过则消息重回队列
            _ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
            return
        try:
            ark_resp = call_ark(req)  # 调用方舟大模型
            print(f"请求:{req[:15]} | 响应:{ark_resp[:20]}")
            _ch.basic_ack(delivery_tag=method.delivery_tag)  # 消费成功手动ACK
        except Exception as e:
            _ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)  # 失败重试

    ch.basic_consume(queue=MQ_QUEUE, on_message_callback=callback, auto_ack=False)
    print(f"消费启动｜初始QPS:{FLOW['curr_qps']} 最大QPS:{FLOW['max_qps']} 斜率可控")
    ch.start_consuming()

# ===================== 启动入口 =====================
if __name__ == '__main__':
    # 启动流控后台线程
    threading.Thread(target=flow_schedule, daemon=True).start()
    threading.Thread(target=reset_count, daemon=True).start()

    # 方式1: 模拟Client端峰值流量【生产端】- 按需执行
    # for i in range(100): producer({"msg": f"大模型请求{i}: 技术方案总结"})

    # 方式2: 启动消费端【填谷+控速+调用方舟】- 主运行
    consumer()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354package main

import (
  "bytes"
  "encoding/json"
  "flag"
  "fmt"
  "io"
  "log"
  "net/http"
  "os"
  "sync"
  "time"

  amqp "github.com/rabbitmq/amqp091-go"
)

// ===================== 全局配置 =====================
const (
  MQQueue    = "client_ark_queue"
  MQUrl      = "amqp://guest:guest@127.0.0.1:5672/"
  ArkModel   = "doubao-seed-2-1-pro-260915"
  ArkBaseUrl = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
)

var (
  // Get API Key: https://ark.volcengine.com/region:cn-beijing/apikey
  ArkApiKey = os.Getenv("ARK_API_KEY")
  MockMode  = false
)

// FlowConfig 流控配置
type FlowConfig struct {
  CurrQPS  int
  MaxQPS   int
  Step     int
  Interval time.Duration
  Count    int
  mu       sync.Mutex
}

// 全局流控对象
var flow = &FlowConfig{
  CurrQPS:  5,
  MaxQPS:   30,
  Step:     2,
  Interval: 5 * time.Second,
  Count:    0,
}

// ===================== 1. 极简斜率可控流控【核心】=====================

// startFlowSchedule 后台协程：匀速增长消费速率，控制斜率
func startFlowSchedule() {
  ticker := time.NewTicker(flow.Interval)
  defer ticker.Stop()

  for range ticker.C {
    flow.mu.Lock()
    if flow.CurrQPS < flow.MaxQPS {
      flow.CurrQPS += flow.Step
      // 确保不超过最大值
      if flow.CurrQPS > flow.MaxQPS {
        flow.CurrQPS = flow.MaxQPS
      }
      log.Printf("QPS Limit Increased: Current=%d, Max=%d", flow.CurrQPS, flow.MaxQPS)
    }
    flow.mu.Unlock()
  }
}

// startResetCount 后台协程：每秒重置请求计数器，精准控QPS
func startResetCount() {
  ticker := time.NewTicker(1 * time.Second)
  defer ticker.Stop()

  for range ticker.C {
    flow.mu.Lock()
    flow.Count = 0
    flow.mu.Unlock()
  }
}

// allow 流控校验：是否允许发起本次请求
func allow() bool {
  flow.mu.Lock()
  defer flow.mu.Unlock()

  if flow.Count < flow.CurrQPS {
    flow.Count++
    return true
  }
  return false
}

// ===================== 3. 方舟API调用【大模型请求核心】=====================

type ArkRequest struct {
  Model    string    `json:"model"`
  Messages []Message `json:"messages"`
}

type Message struct {
  Role    string `json:"role"`
  Content string `json:"content"`
}

type ArkResponse struct {
  Choices []struct {
    Message Message `json:"message"`
  } `json:"choices"`
  Error *struct {
    Message string `json:"message"`
    Type    string `json:"type"`
  } `json:"error,omitempty"`
}

// callArk 调用方舟大模型API
func callArk(msg string) (string, error) {
  if MockMode {
    time.Sleep(50 * time.Millisecond) // 模拟网络延迟
    return "Mocked Response: " + msg, nil
  }

  reqBody := ArkRequest{
    Model: ArkModel,
    Messages: []Message{
      {Role: "user", Content: msg},
    },
  }

  jsonData, err := json.Marshal(reqBody)
  if err != nil {
    return "", err
  }

  req, err := http.NewRequest("POST", ArkBaseUrl, bytes.NewBuffer(jsonData))
  if err != nil {
    return "", err
  }

  req.Header.Set("Content-Type", "application/json")
  req.Header.Set("Authorization", "Bearer "+ArkApiKey)

  client := &http.Client{Timeout: 60 * time.Second}
  resp, err := client.Do(req)
  if err != nil {
    return "", err
  }
  defer resp.Body.Close()

  body, err := io.ReadAll(resp.Body)
  if err != nil {
    return "", err
  }

  if resp.StatusCode != http.StatusOK {
    return "", fmt.Errorf("API error: status=%d, body=%s", resp.StatusCode, string(body))
  }

  var arkResp ArkResponse
  if err := json.Unmarshal(body, &arkResp); err != nil {
    return "", err
  }

  if arkResp.Error != nil {
    return "", fmt.Errorf("API error: %s", arkResp.Error.Message)
  }

  if len(arkResp.Choices) > 0 {
    return arkResp.Choices[0].Message.Content, nil
  }
  return "", fmt.Errorf("empty choices in response")
}

// ===================== 4. MQ消费者【填谷+斜率控速+调用方舟，核心主逻辑】=====================

// startConsumer 消费端核心：削峰后的流量，斜率可控匀速消费，调用方舟大模型
func startConsumer() {
  conn, err := amqp.Dial(MQUrl)
  if err != nil {
    log.Fatalf("Failed to connect to RabbitMQ: %v", err)
  }
  defer conn.Close()

  ch, err := conn.Channel()
  if err != nil {
    log.Fatalf("Failed to open a channel: %v", err)
  }
  defer ch.Close()

  // 声明队列 (Durable=true)
  q, err := ch.QueueDeclare(
    MQQueue, // name
    true,    // durable
    false,   // delete when unused
    false,   // exclusive
    false,   // no-wait
    nil,     // arguments
  )
  if err != nil {
    log.Fatalf("Failed to declare a queue: %v", err)
  }

  // 设置QoS (PrefetchCount=1)
  err = ch.Qos(
    1,     // prefetch count
    0,     // prefetch size
    false, // global
  )
  if err != nil {
    log.Fatalf("Failed to set QoS: %v", err)
  }

  msgs, err := ch.Consume(
    q.Name, // queue
    "",     // consumer
    false,  // auto-ack (设置为false，手动ACK)
    false,  // exclusive
    false,  // no-local
    false,  // no-wait
    nil,    // args
  )
  if err != nil {
    log.Fatalf("Failed to register a consumer: %v", err)
  }

  log.Printf("消费启动｜初始QPS:%d 最大QPS:%d 斜率可控", flow.CurrQPS, flow.MaxQPS)

  forever := make(chan bool)

  go func() {
    for d := range msgs {
      // 解析消息
      var reqData map[string]string
      if err := json.Unmarshal(d.Body, &reqData); err != nil {
        log.Printf("Error parsing JSON: %v", err)
        d.Nack(false, false) // 无法解析，丢弃或放入死信队列
        continue
      }
      msgContent := reqData["msg"]

      // 流控校验
      if !allow() {
        // 不通过则消息重回队列
        d.Nack(false, true)
        // 稍微sleep一下避免空转太快
        time.Sleep(10 * time.Millisecond)
        continue
      }

      // 调用方舟大模型
      respContent, err := callArk(msgContent)
      if err != nil {
        log.Printf("调用方舟失败: %v", err)
        d.Nack(false, true) // 失败重试
      } else {
        // 打印日志 (截取部分长度)
        reqSnippet := msgContent
        if len(reqSnippet) > 15 {
          reqSnippet = reqSnippet[:15]
        }
        respSnippet := respContent
        if len(respSnippet) > 20 {
          respSnippet = respSnippet[:20]
        }
        log.Printf("请求:%s | 响应:%s", reqSnippet, respSnippet)


        // 消费成功手动ACK
        d.Ack(false)
      }
    }
  }()

  <-forever
}

// ===================== 2. MQ生产者【Client端请求投递，削峰】=====================

// startProducer 投递Client请求到MQ，无速率限制，承接峰值流量
func startProducer(count int) {
  conn, err := amqp.Dial(MQUrl)
  if err != nil {
    log.Fatalf("Failed to connect to RabbitMQ: %v", err)
  }
  defer conn.Close()

  ch, err := conn.Channel()
  if err != nil {
    log.Fatalf("Failed to open a channel: %v", err)
  }
  defer ch.Close()

  q, err := ch.QueueDeclare(
    MQQueue, // name
    true,    // durable
    false,   // delete when unused
    false,   // exclusive
    false,   // no-wait
    nil,     // arguments
  )
  if err != nil {
    log.Fatalf("Failed to declare a queue: %v", err)
  }

  log.Printf("开始投递 %d 条消息...", count)
  for i := 0; i < count; i++ {
    msgContent := fmt.Sprintf("大模型请求%d: 技术方案总结", i)
    reqData := map[string]string{"msg": msgContent}
    body, _ := json.Marshal(reqData)

    err = ch.Publish(
      "",     // exchange
      q.Name, // routing key
      false,  // mandatory
      false,  // immediate
      amqp.Publishing{
        DeliveryMode: amqp.Persistent,
        ContentType:  "application/json",
        Body:         body,
      })
    if err != nil {
      log.Fatalf("Failed to publish a message: %v", err)
    }
  }
  log.Printf("成功投递 %d 条消息", count)
}

// ===================== 启动入口 =====================

func main() {
  mode := flag.String("mode", "consumer", "运行模式: consumer | producer")
  count := flag.Int("count", 100, "生产者发送消息数量")
  mock := flag.Bool("mock", false, "是否开启Mock模式(不真实调用API)")
  flag.Parse()

  if *mock {
    MockMode = true
    log.Println("Mock模式已开启")
  }

  if *mode == "producer" {
    startProducer(*count)
    return
  }

  // 启动流控后台协程
  go startFlowSchedule()
  go startResetCount()

  // 启动消费端
  startConsumer()
}