代码收藏家技术教程 2025-01-24

解决Python调用websocket.WebSocket()卡住的问题

提问

Python调用websocket.WebSocket()卡住的问题'

解答

代码卡住的问题可能由几个方面引起，包括网络延迟、服务器响应超时、WebSocket连接问题或是代码逻辑中存在的阻塞点。为了解决或缓解这些问题，可以考虑以下几个方面的优化：

增加超时设置：urllib.request.urlopen函数默认可能没有设置超时或设置得较长，可以通过添加timeout参数来设置一个合理的超时时间。例如，设置为10秒：
```
response = urllib.request.urlopen(req, timeout=10)
```
异常处理：在网络请求周围添加try-except块，以捕获超时或其他网络相关的异常。这样可以在请求失败时进行适当的错误处理，而不是让整个程序卡住。
```
try:
    response = urllib.request.urlopen(req, timeout=10)
except Exception as e:
    print(f"请求失败: {e}")
    return None
```
WebSocket消息处理优化：get_images函数中，ws.recv()是一个阻塞调用，它会一直等待直到收到新消息。如果服务器停止发送消息，这将导致无限等待。可以设置一个超时时间，避免无限等待：
```
ws.settimeout(10)  # 设置WebSocket超时时间
try:
    out = ws.recv()
except websocket.WebSocketTimeoutException:
    print("WebSocket接收消息超时")
    break
```
验证服务器响应：确保从服务器接收到的数据是有效的，特别是在处理queue_prompt和get_history的响应时。如果服务器返回的是错误或异常信息，应该对其进行处理。
限制并发数量：如果同时处理多个请求，考虑限制并发的数量。太多并发请求可能会导致服务器响应变慢，甚至导致客户端处理不过来。可以使用线程池或协程来控制并发量。
日志记录：增加更多的日志记录，帮助定位是哪一部分的请求或处理导致了程序卡住。记录每个关键步骤的开始和完成，可以帮助快速定位问题所在。

通过以上的优化，可以提高程序的健壮性和响应性，减少因网络问题或服务器响应延迟导致的卡顿现象。

以下是一个ComfyUI调用解决卡顿的一个实例。

增加请求超时处理：在urllib.request.urlopen调用中增加超时设置，并对可能出现的异常进行捕获处理。
WebSocket接收消息时增加超时机制：通过在WebSocket连接前设置超时，并在接收消息时捕获超时异常。
增加异常处理逻辑：对可能发生的异常情况进行捕获和处理，确保程序能够在遇到错误时优雅地终止或重试，而不是无限等待或崩溃。

# 导入所需库
import websocket
import uuid
import json
import urllib.request
import urllib.parse
import random
import pandas as pd
import ssl

# 设置服务器地址和客户端ID
server_address = "127.0.0.1:8188"
client_id = str(uuid.uuid4())

# 增加超时设置
def safe_urlopen(url, data=None, timeout=10):
    try:
        response = urllib.request.urlopen(url, data=data, timeout=timeout)
        return response
    except Exception as e:
        print(f"网络请求出错: {e}")
        return None

# 定义向服务器发送提示的函数，增加了超时处理
def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request("http://{}/prompt".format(server_address), data=data)
    response = safe_urlopen(req)
    if response:
        return json.loads(response.read())
    else:
        return None

# 修改WebSocket消息接收逻辑，增加超时机制
def get_images(ws, prompt, timeout=10):
    prompt_response = queue_prompt(prompt)
    if not prompt_response:
        print("无法发送提示到服务器")
        return {}
    prompt_id = prompt_response['prompt_id']
    output_images = {}
    ws.settimeout(timeout)  # 设置WebSocket的超时
    try:
        while True:
            out = ws.recv()
            if isinstance(out, str):
                message = json.loads(out)
                if message['type'] == 'executing':
                    data = message['data']
                    if data['node'] is None and data['prompt_id'] == prompt_id:
                        break  # 执行完成
    except websocket.WebSocketTimeoutException:
        print("WebSocket接收消息超时")
        return output_images  # 返回当前收集到的图像数据

    # 后续处理逻辑...

这些修改是基于您原始代码片段的直接修改，可能需要根据实际情况进行调整。特别是错误处理和异常处理的逻辑，应根据您的具体需求和预期行为进行定制。

# 继续定义从服务器下载图像数据的函数，增加了超时处理
def get_image(filename, subfolder, folder_type, timeout=10):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    response = safe_urlopen("http://{}/view?{}".format(server_address, url_values), timeout=timeout)
    if response:
        return response.read()
    else:
        return None

# 继续定义获取历史记录的函数，增加了超时处理
def get_history(prompt_id, timeout=10):
    response = safe_urlopen("http://{}/history/{}".format(server_address, prompt_id), timeout=timeout)
    if response:
        return json.loads(response.read())
    else:
        return None

# 调整和优化get_images函数
def get_images(ws, prompt, timeout=10):
    prompt_response = queue_prompt(prompt)
    if not prompt_response:
        print("无法发送提示到服务器")
        return {}
    prompt_id = prompt_response.get('prompt_id', '')
    output_images = {}
    ws.settimeout(timeout)
    try:
        while True:
            out = ws.recv()
            if out:
                message = json.loads(out)
                if message['type'] == 'executing' and message.get('data', {}).get('prompt_id', '') == prompt_id:
                    break  # 如果执行完成，则退出循环
    except websocket.WebSocketTimeoutException:
        print("WebSocket接收消息超时")
    except Exception as e:
        print(f"WebSocket接收消息异常: {e}")

    history = get_history(prompt_id)
    if history:
        history = history.get(prompt_id, {})
        for node_id, node_output in history.get('outputs', {}).items():
            if 'images' in node_output:
                images_output = []
                for image in node_output['images']:
                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
                    if image_data:
                        images_output.append(image_data)
                output_images[node_id] = images_output

    return output_images

# 注意：以下是主程序部分，确保在主程序逻辑中正确处理异常和超时情况
if __name__ == "__main__":
    # 省略示例JSON字符串和DataFrame处理部分

    # 创建WebSocket连接
    try:
        ws = websocket.WebSocket()
        ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id), timeout=10)
        images = get_images(ws, prompt)
        print("图像处理完成")
    except Exception as e:
        print(f"创建WebSocket连接或处理图像时发生错误: {e}")
    finally:
        ws.close()

这里引入了几个关键的改进：