深入解析openai库流式输出:确保网络连接安全关闭的关键技巧

279 阅读7分钟

本篇文章首发于【算法工程笔记】,更多内容,欢迎关注。

注:本文章提到的 openai ,均指 python 语言的一个库OpenAI发明 GPT 系列模型的公司。

在上一篇文章——《使用openai库进行流式输出时,到底发生了什么》中,我们着重介绍了 openai 库流式输出的基本原理,以及当使用这个库的流式输出功能时,具体的网络请求是什么

在这篇文章,我们着重分析几种常见流式请求时,网络的连接情况

文章内容有点长,如果没有时间读完的话,请牢记下面这段话

客户端未完整遍历流式响应时,连接不会自动断开,会长期停留在 CLOSE_WAIT 状态,可能导致连接资源耗尽
显式调用 response.close() 是确保连接安全关闭的关键

在进行异常情况分析前,我们先介绍下 TCP 连接中的几种状态。

TCP连接中,客户端和服务端的状态

关于这部分内容,网上已经有很多资料,这里就采用 《实战!我用 Wireshark 让你“看见“ TCP》 中的一张图片:

19.jpg

TCP 的握手和挥手过程

在此,只介绍其中几个重要的状态:

  • ESTABLISHED表示 TCP 连接已经成功建立
  • CLOSE_WAIT:表明本方已经接收到另一方发来的关闭请求(FIN报文),但还没有完全关闭连接,需要等待本方去关闭连接
  • TIME_WAIT表示收到了对方的FIN报文,并发送出了ACK报文TIME_WAIT状态下的TCP连接会等待2*MSL(Max Segment Lifetime,最大分段生存期,指一个TCP报文在Internet上的最长生存时间),然后自行切换到 CLOSED 状态

明确了 TCP 连接中的状态,接下来就分析几种典型流式输出情况。

不同流式输出情况下,网络连接情况分析

这里的主体代码与上一篇文章中的类似,只是增加了调用次数,以便查看具体的连接切换状态:

import os
from openai import OpenAI

base_url, api_key, model = BASE_URL, API_KEY, MODEL
client = OpenAI(api_key=api_key, base_url=base_url)

def test(stream=True):
    query = f"please introduce yourself briefly, no more than ten words."
    messages = [{"role""user""content": query}]
    req_dic = {"model": model, "messages": messages, "stream": stream}
    response = client.chat.completions.create(**req_dic)
    time.sleep(5)
    if stream:
        res = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                res += chunk.choices[0].delta.content
        pass

import time
test()
time.sleep(10)
test()
time.sleep(10)

客户端完整遍历流式结果

这种情况对应流式输出的绝大多数情况,即调用流式接口后,遍历流式输出结果。

通过抓取网络包和观察对应的网络连接状态,可以看到客户端的连接过程。

这里观察对应网络状态是在客户端执行 netstat 命令,具体如下:

watch -n 1 " netstat -nltpa | grep '8700' "

这里 netstat -nltpa 主要作用是显示当前的网络连接信息。其中:
-n:以数字形式显示地址和端口,避免DNS查询
-l:显示正在监听(LISTEN)的端口
-t:仅显示 TCP 连接
-p:显示与连接对应的进程ID(PID)和进程名称
-a:显示所有连接状态,包括已建立和监听状态

运行上述代码后,可以观察到的网络连接状态的变更如下:

ESTABLISHED -> CLOSE_WAIT -> CLOSE ESTABLISHED -> CLOSE_WAIT -> CLOSE

wireshark 中显示的网络数据包传输图如下(仅显示两次请求之间的过程,下同):

数据包传输图一

数据包传输图一

可以看到以下现象:

  1. 服务端收到请求后会持续向客户端发送响应,直至响应结束,此时客户端一直处于 ESTABLISHED 状态(第40、41帧)
  2. 客户端遍历完响应后,服务端发送 FIN 消息,发起挥手消息(第42帧)
  3. 客户端收到挥手消息,发送确认消息,同时进入 CLOSE_WAIT 状态(第43帧)
  4. 客户端再次请求前,发送挥手消息,断开上一个连接,此时连接进入 CLOSE 状态(第44帧)
  5. 客户端发送第二次请求(第45帧)

对应的客户端连接的状态变化如下图:

连接状态转换.png

客户端完整遍历流式结果后,发起close请求

与第一种情况不同的是,客户端在遍历完响应后,通过 response.close() 主动发起 close 请求

脚本上与主体代码的主要区别是在 for 循环后增加如下代码

        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                res += chunk.choices[0].delta.content
        pass
        time.sleep(2)
        response.close()

网络连接状态的变更如下:

ESTABLISHED -> CLOSE_WAIT -> CLOSE
ESTABLISHED -> CLOSE_WAIT -> CLOSE

wireshark 中显示的网络数据包传输图如下:

数据包传输图二

数据包传输图二

可以看到,这种操作下的连接情况与第一种完全一样,造成这种情况的原因我们可以在 openai 包中找到——在流式响应被遍历之后,会自动调用 close() 方法,具体可以参考 openai 包的 openai._streaming.Stream.close()httpx 包的httpx._model.Response.iter_raw()

客户端未完整遍历流式结果

与第一种情况不同的是,客户端只发送请求,不对返回的响应进行遍历或者仅遍历某一部分的结果,这实际上可以对应以下的实际情况:

  1. 当输出的结果包含某些敏感词时,不再对结果进行遍历
  2. 用户主动点击停止生成

脚本上与主体代码的主要区别是在 for 循环中增加退出循环的代码

        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                res += chunk.choices[0].delta.content
                if len(res) > 2:
                    break
        pass

网络连接状态的变更如下:

ESTABLISHED -> CLOSE_WAIT
ESTABLISHED -> CLOSE_WAIT

注意,这里的两个连接没有到最终的 CLOSE 状态!!!

与此同时,wireshark 中显示的网络数据包传输图如下:

数据包传输图三

数据包传输图三

可以看到,在客户端没有完整遍历流式响应的情况下,客户端与服务端的连接是没有断开的!!!

连接一直处于CLOSE_WAIT的状态,而当这种状态的连接数量过多时,可能会导致客户端无法与服务端连接

因此需要在实际生产环境中避免这种状态!!!

具体如何避免,请看下一小结。

客户端未完整遍历流式结果、客户端发起close请求

与上一种情况不同,当客户端决定不遍历完响应结果时,也可以主动发起close请求,具体代码如下:

        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                res += chunk.choices[0].delta.content
                if len(res) > 2:
                    response.close()
                    break
        pass

此时,网络连接状态的变更如下:

ESTABLISHED -> TIME_WAIT
ESTABLISHED -> TIME_WAIT

注意,这里的两个连接没有到最终的 CLOSE 状态!!!

与此同时,wireshark 中显示的网络数据包传输图如下:

数据包传输图四

数据包传输图四

可以看到,当客户端调用 response.close() 之后,会主动发送一个 FIN,ACK 的挥手请求,服务端收到后也会回复 ACK 和 FIN ,客户端收到后返回 ACK 消息,同时连接进入 TIME_WAIT 状态,当连接进入 TIME_WAIT 状态之后,等待2倍的 MSL 时间后,连接就会完全关闭。

25.jpg

TCP四次挥手

因此,当流式输出没有被完全遍历时(结果包含敏感词或者用户主动点击停止),客户端需要主动调用 close 方法,以确保连接正常关闭

以上就是几种常见流式请求时,网络的连接情况分析,感谢阅读到这里,下次再见!