Python 3.10中UDP接收模型权重后Pickle反序列化报错

开展网络组播实验时，服务器需将处理后的模型权重（ndarray元组）回传给客户端。使用pickle.dumps()和pickle.loads()进行序列化与反序列化时，持续报错：_pickle.UnpicklingError : invalid load key, 'XXX'（XXX为dump中的某个键）。当前使用Python 3.10，以下是最小复现代码：

发送端代码 #

import pickle
import socket
import struct
import sys
import numpy as np

group_ip_address = "239.0.0.100"
port = 10000
source_ip = "192.168.68.116"
source_port = 12345

print("Multicast sender transmitting data to ", group_ip_address, "/", port, ". Local address: ", source_ip,
      "/", source_port)

message = tuple(np.random.rand(*(100,100)) for _ in range(10))

multicast_group = (group_ip_address, port)

# Create UDP socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# bind to selected address
sock.bind((source_ip, source_port))

# Set the time-to-live for messages to 1 so they do not go past the
# local network segment.
ttl = struct.pack('b', 1)
sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, ttl)

try:
    BUFFER_SIZE = 4096

    # Send data to the multicast group
    print(sys.getsizeof(message))
    print(message)
    serialized = pickle.dumps(message)
    for i in range(0, len(serialized), BUFFER_SIZE):
        data = serialized[i:i + BUFFER_SIZE]
        sock.sendto(data, multicast_group)

        while True:
            sock.settimeout(1)
            print('waiting to receive acknowledgment')
            try:

                data, server = sock.recvfrom(BUFFER_SIZE)
            except socket.timeout:
                print('timed out, resending package')
                sock.sendto(data, multicast_group)
            else:
                print('received "%s" from %s' % (data, server))
                break
    # tell edge server to stop the recv loop
    sock.sendto(b'END', multicast_group)

finally:
    print('closing socket')
    sock.close()

接收端代码 #

import pickle
import socket
import struct

multicast_group = '239.0.0.100'
port = 10000

print("Multicast receiver listening for data at ", multicast_group, "/", port)

server_address = ('', port)

# Create UDP socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# Bind to the server address
sock.bind(server_address)

# Tell the operating system to add the socket to the multicast group
# on all interfaces.
group = socket.inet_aton(multicast_group)
mreq = struct.pack('4sL', group, socket.INADDR_ANY)
sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)

BUFFER_SIZE = 4096
packets = []
while True:
    data, address = sock.recvfrom(BUFFER_SIZE)
    print(data)
    if not data or data.endswith(b'END'):
        break
    packets.append(data)
    sock.sendto(b'ack', address)

bytes = b''.join(packets)
new_aggregated_parameters = pickle.loads(bytes)
print(type(new_aggregated_parameters))

核心问题分析 #

1. 发送端变量覆盖：发送数据包后，recvfrom将ACK内容存入了和待发送数据包同名的data变量，导致超时重发时实际发送的是ACK内容而非原数据包，最终接收端拼接的数据不完整。
2. 接收端结束信号误判：使用data.endswith(b'END')判断结束，若某个数据包末尾恰好包含b'END'字节序列，会提前终止接收，导致序列化数据不完整。

修改后的代码 #

发送端（修正变量覆盖问题）

import pickle
import socket
import struct
import sys
import numpy as np

group_ip_address = "239.0.0.100"
port = 10000
source_ip = "192.168.68.116"
source_port = 12345

print("Multicast sender transmitting data to ", group_ip_address, "/", port, ". Local address: ", source_ip,
      "/", source_port)

message = tuple(np.random.rand(*(100,100)) for _ in range(10))

multicast_group = (group_ip_address, port)

# Create UDP socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# bind to selected address
sock.bind((source_ip, source_port))

# Set the time-to-live for messages to 1 so they do not go past the
# local network segment.
ttl = struct.pack('b', 1)
sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, ttl)

try:
    BUFFER_SIZE = 4096

    # Send data to the multicast group
    print(sys.getsizeof(message))
    print(message)
    serialized = pickle.dumps(message)
    for i in range(0, len(serialized), BUFFER_SIZE):
        # 重命名变量，避免被ACK内容覆盖
        packet = serialized[i:i + BUFFER_SIZE]
        sock.sendto(packet, multicast_group)

        while True:
            sock.settimeout(1)
            print('waiting to receive acknowledgment')
            try:
                # 使用单独变量存储ACK
                ack_data, server = sock.recvfrom(BUFFER_SIZE)
            except socket.timeout:
                print('timed out, resending package')
                sock.sendto(packet, multicast_group)
            else:
                print('received "%s" from %s' % (ack_data, server))
                break
    # 发送结束信号
    sock.sendto(b'END', multicast_group)

finally:
    print('closing socket')
    sock.close()

接收端（修正结束信号判断）

import pickle
import socket
import struct

multicast_group = '239.0.0.100'
port = 10000

print("Multicast receiver listening for data at ", multicast_group, "/", port)

server_address = ('', port)

# Create UDP socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# Bind to the server address
sock.bind(server_address)

# Tell the operating system to add the socket to the multicast group
# on all interfaces.
group = socket.inet_aton(multicast_group)
mreq = struct.pack('4sL', group, socket.INADDR_ANY)
sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)

BUFFER_SIZE = 4096
packets = []
while True:
    data, address = sock.recvfrom(BUFFER_SIZE)
    print(data)
    # 精确判断结束信号，仅当收到完整的b'END'时终止
    if data == b'END':
        break
    packets.append(data)
    sock.sendto(b'ack', address)

bytes = b''.join(packets)
new_aggregated_parameters = pickle.loads(bytes)
print(type(new_aggregated_parameters))

额外建议 #

• 组播场景若存在多个接收端，ACK机制会引发混乱，建议改用先发送数据包总长度的方式：发送端先把序列化数据的总长度用struct打包发送，接收端先接收长度，再根据长度接收所有数据包，避免依赖ACK。
• 可考虑使用cloudpickle替代pickle，对numpy数组的序列化兼容性更好。

内容的提问来源于stack exchange，提问作者kyro121