feat(lekiwi): stream observations as ZMQ multipart (-25% bandwidth) (#4088)

Observations were sent as base64-in-JSON, which inflates every camera
frame 33% purely to fit binary JPEG inside a text protocol (~11 Mbps at
3 cameras x 30 Hz). Send a ZMQ multipart message instead: a JSON header
frame (state + camera order) followed by one raw JPEG frame per camera.

Benchmarked on hardware: -25% wire size (invariant across contention),
lower and tighter latency, and no dropped frames under load where the
base64 format stalled.

ZMQ_CONFLATE does not support multipart, so the observation sockets use
2-deep high-water marks; the client's existing drain-to-latest loop
preserves the keep-newest behavior.

Breaking wire change: host and client must run the same version.

Co-authored-by: Xingdong Zuo <18168681+zuoxingdong@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
Xingdong Zuo
2026-07-31 00:32:11 +09:00
committed by GitHub
parent 72a1858015
commit 4c302572c0
2 changed files with 46 additions and 41 deletions
+30 -30
View File
@@ -14,7 +14,6 @@
# TODO(aliberts, Steven, Pepijn): use gRPC calls instead of zmq? # TODO(aliberts, Steven, Pepijn): use gRPC calls instead of zmq?
import base64
import json import json
import logging import logging
from functools import cached_property from functools import cached_property
@@ -134,7 +133,9 @@ class LeKiwiClient(Robot):
self.zmq_observation_socket = self.zmq_context.socket(zmq.PULL) self.zmq_observation_socket = self.zmq_context.socket(zmq.PULL)
zmq_observations_locator = f"tcp://{self.remote_ip}:{self.port_zmq_observations}" zmq_observations_locator = f"tcp://{self.remote_ip}:{self.port_zmq_observations}"
self.zmq_observation_socket.connect(zmq_observations_locator) self.zmq_observation_socket.connect(zmq_observations_locator)
self.zmq_observation_socket.setsockopt(zmq.CONFLATE, 1) # CONFLATE does not support multipart messages; a small receive queue plus
# the existing drain-to-latest loop keeps newest-only semantics.
self.zmq_observation_socket.setsockopt(zmq.RCVHWM, 2)
poller = zmq.Poller() poller = zmq.Poller()
poller.register(self.zmq_observation_socket, zmq.POLLIN) poller.register(self.zmq_observation_socket, zmq.POLLIN)
@@ -147,8 +148,8 @@ class LeKiwiClient(Robot):
def calibrate(self) -> None: def calibrate(self) -> None:
pass pass
def _poll_and_get_latest_message(self) -> str | None: def _poll_and_get_latest_message(self) -> list[bytes] | None:
"""Polls the ZMQ socket for a limited time and returns the latest message string.""" """Polls the ZMQ socket for a limited time and returns the latest message's frames."""
zmq = self._zmq zmq = self._zmq
poller = zmq.Poller() poller = zmq.Poller()
poller.register(self.zmq_observation_socket, zmq.POLLIN) poller.register(self.zmq_observation_socket, zmq.POLLIN)
@@ -166,7 +167,7 @@ class LeKiwiClient(Robot):
last_msg = None last_msg = None
while True: while True:
try: try:
msg = self.zmq_observation_socket.recv_string(zmq.NOBLOCK) msg = self.zmq_observation_socket.recv_multipart(zmq.NOBLOCK)
last_msg = msg last_msg = msg
except zmq.Again: except zmq.Again:
break break
@@ -176,28 +177,27 @@ class LeKiwiClient(Robot):
return last_msg return last_msg
def _parse_observation_json(self, obs_string: str) -> RobotObservation | None: def _parse_observation(self, frames: list[bytes]) -> RobotObservation | None:
"""Parses the JSON observation string.""" """Parses a multipart observation: JSON header + one raw JPEG frame per camera."""
try: try:
return json.loads(obs_string) header = json.loads(frames[0])
except json.JSONDecodeError as e: cam_names = header.pop("_cams")
logging.error(f"Error decoding JSON observation: {e}") observation: RobotObservation = header
for cam_name, jpeg in zip(cam_names, frames[1:], strict=True):
observation[cam_name] = jpeg
return observation
except (json.JSONDecodeError, KeyError, ValueError) as e:
logging.error(f"Error decoding observation: {e}")
return None return None
def _decode_image_from_b64(self, image_b64: str) -> np.ndarray | None: def _decode_image(self, jpeg: bytes) -> np.ndarray | None:
"""Decodes a base64 encoded image string to an OpenCV image.""" """Decodes a raw JPEG buffer to an OpenCV image."""
if not image_b64: if not jpeg:
return None
try:
jpg_data = base64.b64decode(image_b64)
np_arr = np.frombuffer(jpg_data, dtype=np.uint8)
frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
if frame is None:
logging.warning("cv2.imdecode returned None for an image.")
return frame
except (TypeError, ValueError) as e:
logging.error(f"Error decoding base64 image data: {e}")
return None return None
frame = cv2.imdecode(np.frombuffer(jpeg, dtype=np.uint8), cv2.IMREAD_COLOR)
if frame is None:
logging.warning("cv2.imdecode returned None for an image.")
return frame
def _remote_state_from_obs( def _remote_state_from_obs(
self, observation: RobotObservation self, observation: RobotObservation
@@ -212,10 +212,10 @@ class LeKiwiClient(Robot):
# Decode images # Decode images
current_frames: dict[str, np.ndarray] = {} current_frames: dict[str, np.ndarray] = {}
for cam_name, image_b64 in observation.items(): for cam_name, jpeg in observation.items():
if cam_name not in self._cameras_ft: if cam_name not in self._cameras_ft:
continue continue
frame = self._decode_image_from_b64(image_b64) frame = self._decode_image(jpeg)
if frame is not None: if frame is not None:
current_frames[cam_name] = frame current_frames[cam_name] = frame
@@ -230,15 +230,15 @@ class LeKiwiClient(Robot):
If no new data arrives or decoding fails, returns the last known values. If no new data arrives or decoding fails, returns the last known values.
""" """
# 1. Get the latest message string from the socket # 1. Get the latest message's frames from the socket
latest_message_str = self._poll_and_get_latest_message() latest_frames = self._poll_and_get_latest_message()
# 2. If no message, return cached data # 2. If no message, return cached data
if latest_message_str is None: if latest_frames is None:
return self.last_frames, self.last_remote_state return self.last_frames, self.last_remote_state
# 3. Parse the JSON message # 3. Parse the multipart message
observation = self._parse_observation_json(latest_message_str) observation = self._parse_observation(latest_frames)
# 4. If JSON parsing failed, return cached data # 4. If JSON parsing failed, return cached data
if observation is None: if observation is None:
+16 -11
View File
@@ -14,7 +14,6 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import base64
import json import json
import logging import logging
import time import time
@@ -44,7 +43,9 @@ class LeKiwiHost:
self.zmq_cmd_socket.bind(f"tcp://*:{config.port_zmq_cmd}") self.zmq_cmd_socket.bind(f"tcp://*:{config.port_zmq_cmd}")
self.zmq_observation_socket = self.zmq_context.socket(zmq.PUSH) self.zmq_observation_socket = self.zmq_context.socket(zmq.PUSH)
self.zmq_observation_socket.setsockopt(zmq.CONFLATE, 1) # CONFLATE does not support multipart messages; a 2-deep send queue keeps
# near-latest-only semantics and sheds stale observations during stalls.
self.zmq_observation_socket.setsockopt(zmq.SNDHWM, 2)
self.zmq_observation_socket.bind(f"tcp://*:{config.port_zmq_observations}") self.zmq_observation_socket.bind(f"tcp://*:{config.port_zmq_observations}")
self.connection_time_s = config.connection_time_s self.connection_time_s = config.connection_time_s
@@ -99,19 +100,23 @@ def main(cfg: LeKiwiServerConfig):
last_observation = robot.get_observation() last_observation = robot.get_observation()
# Encode ndarrays to base64 strings # Send one multipart message: a JSON header frame (state + camera
for cam_key, _ in robot.cameras.items(): # order) followed by one raw JPEG frame per camera. Raw JPEG avoids
ret, buffer = cv2.imencode( # the 33% base64 inflation of embedding binary data in JSON.
".jpg", last_observation[cam_key], [int(cv2.IMWRITE_JPEG_QUALITY), 90] cam_keys = list(robot.cameras.keys())
jpeg_frames = []
for cam_key in cam_keys:
ret, jpeg = cv2.imencode(
".jpg", last_observation.pop(cam_key), [int(cv2.IMWRITE_JPEG_QUALITY), 90]
) )
if ret: jpeg_frames.append(jpeg if ret else b"")
last_observation[cam_key] = base64.b64encode(buffer).decode("utf-8") header = {"_cams": cam_keys, **last_observation}
else:
last_observation[cam_key] = ""
# Send the observation to the remote agent # Send the observation to the remote agent
try: try:
host.zmq_observation_socket.send_string(json.dumps(last_observation), flags=zmq.NOBLOCK) host.zmq_observation_socket.send_multipart(
[json.dumps(header).encode()] + jpeg_frames, flags=zmq.NOBLOCK
)
except zmq.Again: except zmq.Again:
logging.info("Dropping observation, no client connected") logging.info("Dropping observation, no client connected")