[diffusion]: fix scheduler crash on ZMQ messages with unexpected frame counts (#17890)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
Ratish P
2026-02-17 07:15:05 +05:30
committed by GitHub
parent f9c3def7fe
commit 50ca24aebb

View File

@@ -241,9 +241,13 @@ class Scheduler:
if self.receiver is not None:
try:
try:
identity, _, payload = self.receiver.recv_multipart(zmq.NOBLOCK)
recv_reqs = pickle.loads(payload)
except zmq.Again:
# Accept valid REQ envelopes only, ignore malformed/probe frames.
parts = self.receiver.recv_multipart(zmq.NOBLOCK)
identity, payload = parts[0], parts[-1]
# Ignore malformed probes or non-pickle data
recv_reqs = pickle.loads(payload) if len(parts) > 2 else []
except (zmq.Again, pickle.UnpicklingError, IndexError, EOFError):
recv_reqs = []
except zmq.ZMQError:
# re-raise or handle appropriately to let the outer loop continue