Files
frigate/frigate/detectors/plugins/zmq_ipc.py
T
61fd753aed Add accelerated input transform support to zmq detector protocol (#24731)
* Let the zmq detector skip the input transform when the detector applies it

_transform_input converts the frame to the model's layout and dtype before
the detector sees it. For a remote detector that can apply the transform
itself, often on a GPU, that work is done twice over and the tensor is sent
in its most expensive form.

The detector now reports input_transform_offload in the model handshake. When
it does, Frigate sends the frame as captured (uint8 NHWC) and adds the model's
expected input_tensor and input_dtype to the request header, so the detector
can apply them. The model config is unchanged either way: the expected format
is relayed, not overridden.

The choice is read per frame rather than cached, because the detector redoes
its handshake on a timeout or socket error, so restarting or upgrading it
while Frigate runs already works and keeps working here.

Negotiated rather than configured, for two reasons. The detector is what knows
whether it can apply the transform, and a user who merely believed it could
would send unconverted frames to a model that cannot take them. It also needs
no config surface, which the device string format has no room for.

On a 640x640 region, measured on a Ryzen 9 5900XT, the host-side cost of
handing the detector a frame drops from about 4 ms to under 0.1 ms per
detection, and the payload from 4.9 MB to 1.2 MB. Nearly all of it is
tobytes(): np.transpose returns a view and astype keeps its strides by
default, so the transpose is deferred to serialization, where it becomes a
strided gather over the whole float32 tensor. The YUV to RGB conversion is
unaffected and stays on the host.

* Test that the handshake decides whether the input transform is applied

Parameterized over the three handshake values against the header the detector
ends up sending: reported, the frame goes over as captured uint8 NHWC with the
model's input_tensor and input_dtype declared; not reported or absent, the
existing transform runs and those fields are omitted.

A second case drives a real reconnect -- a timed out request, then a handshake
from an upgraded detector -- and checks the next frame follows the new answer
rather than the one from startup.

---------

Co-authored-by: Eric <eric@localhost>
2026-10-23 06:26:47 -06:00

369 lines
14 KiB
Python

import json
import logging
import os
from typing import Any, ClassVar, Literal
import numpy as np
import zmq
from pydantic import ConfigDict, Field
from frigate.detectors.detection_api import DetectionApi
from frigate.detectors.detector_config import BaseDetectorConfig
logger = logging.getLogger(__name__)
DETECTOR_KEY = "zmq"
class ZmqDetectorConfig(BaseDetectorConfig):
"""ZMQ IPC detector that offloads inference to an external process via a ZeroMQ IPC endpoint."""
model_config = ConfigDict(
title="ZMQ IPC",
)
device_spec_field: ClassVar[str] = "endpoint"
type: Literal[DETECTOR_KEY]
endpoint: str = Field(
default="ipc:///tmp/cache/zmq_detector",
title="ZMQ IPC endpoint",
description="The ZMQ endpoint to connect to.",
)
request_timeout_ms: int = Field(
default=200,
title="ZMQ request timeout in milliseconds",
description="Timeout for ZMQ requests in milliseconds.",
)
linger_ms: int = Field(
default=0,
title="ZMQ socket linger in milliseconds",
description="Socket linger period in milliseconds.",
)
class ZmqIpcDetector(DetectionApi):
"""
ZMQ-based detector plugin using a REQ/REP socket over an IPC endpoint.
Protocol:
- Request is sent as a multipart message:
[ header_json_bytes, tensor_bytes ]
where header is a JSON object containing:
{
"shape": List[int],
"dtype": str, # numpy dtype string, e.g. "uint8", "float32"
}
and, when the detector reported input_transform_offload, also:
{
"input_tensor": str, # InputTensorEnum value, e.g. "nhwc", "nchw"
"input_dtype": str, # InputDTypeEnum value: "float", "float_denorm", "int"
}
which describe the transform the detector is expected to apply itself.
tensor_bytes are the raw bytes of the numpy array in C-order.
- Response is expected to be either:
a) Multipart [ header_json_bytes, tensor_bytes ] with header specifying
shape [20,6] and dtype "float32"; or
b) Single frame tensor_bytes of length 20*6*4 bytes (float32).
On any error or timeout, this detector returns a zero array of shape (20, 6).
Model Management:
- On initialization, sends model request to check if model is available
- If model not available, sends model data via ZMQ
- Only starts inference after model is ready
"""
type_key = DETECTOR_KEY
def __init__(self, detector_config: ZmqDetectorConfig):
super().__init__(detector_config)
self._context = zmq.Context()
self._endpoint = detector_config.endpoint
self._request_timeout_ms = detector_config.request_timeout_ms
self._linger_ms = detector_config.linger_ms
# Reported by the detector during the model handshake.
self.input_transform_offload = False
self._socket = None
self._create_socket()
# Model management
self._model_ready = False
self._model_name = self._get_model_name()
# Initialize model if needed
self._initialize_model()
# Preallocate zero result for error paths
self._zero_result = np.zeros((20, 6), np.float32)
def _create_socket(self) -> None:
if self._socket is not None:
try:
self._socket.close(linger=self._linger_ms)
except Exception:
pass
self._socket = self._context.socket(zmq.REQ)
# Apply timeouts and linger so calls don't block indefinitely
self._socket.setsockopt(zmq.RCVTIMEO, self._request_timeout_ms)
self._socket.setsockopt(zmq.SNDTIMEO, self._request_timeout_ms)
self._socket.setsockopt(zmq.LINGER, self._linger_ms)
logger.debug(f"ZMQ detector connecting to {self._endpoint}")
self._socket.connect(self._endpoint)
def _get_model_name(self) -> str:
"""Get the model filename from the detector config."""
model_path = self.detector_config.model.path
return os.path.basename(model_path)
def _initialize_model(self) -> None:
"""Initialize the model by checking availability and transferring if needed."""
try:
logger.info(f"Initializing model: {self._model_name}")
# Check if model is available and transfer if needed
if self._check_and_transfer_model():
logger.info(f"Model {self._model_name} is ready")
self._model_ready = True
else:
logger.error(f"Failed to initialize model {self._model_name}")
except Exception as e:
logger.error(f"Failed to initialize model: {e}")
def _check_and_transfer_model(self) -> bool:
"""Check if model is available and transfer if needed in one atomic operation."""
try:
# Send model availability request
header = {"model_request": True, "model_name": self._model_name}
header_bytes = json.dumps(header).encode("utf-8")
self._socket.send_multipart([header_bytes])
# Temporarily increase timeout for model operations
original_timeout = self._socket.getsockopt(zmq.RCVTIMEO)
self._socket.setsockopt(zmq.RCVTIMEO, 30000)
try:
response_frames = self._socket.recv_multipart()
finally:
self._socket.setsockopt(zmq.RCVTIMEO, original_timeout)
if len(response_frames) == 1:
try:
response = json.loads(response_frames[0].decode("utf-8"))
model_available = response.get("model_available", False)
model_loaded = response.get("model_loaded", False)
self._read_offload_support(response)
if model_available and model_loaded:
return True
elif model_available and not model_loaded:
logger.error("Model exists but failed to load")
return False
else:
return self._send_model_data()
except json.JSONDecodeError:
logger.warning(
"Received non-JSON response for model availability check"
)
return False
else:
logger.warning(
"Received unexpected response format for model availability check"
)
return False
except Exception as e:
logger.error(f"Failed to check and transfer model: {e}")
return False
def _check_model_availability(self) -> bool:
"""Check if the model is available on the detector."""
try:
# Send model availability request
header = {"model_request": True, "model_name": self._model_name}
header_bytes = json.dumps(header).encode("utf-8")
self._socket.send_multipart([header_bytes])
# Receive response
response_frames = self._socket.recv_multipart()
# Check if this is a JSON response (model management)
if len(response_frames) == 1:
try:
response = json.loads(response_frames[0].decode("utf-8"))
model_available = response.get("model_available", False)
model_loaded = response.get("model_loaded", False)
logger.debug(
f"Model availability check: available={model_available}, loaded={model_loaded}"
)
return model_available and model_loaded
except json.JSONDecodeError:
logger.warning(
"Received non-JSON response for model availability check"
)
return False
else:
logger.warning(
"Received unexpected response format for model availability check"
)
return False
except Exception as e:
logger.error(f"Failed to check model availability: {e}")
return False
def _send_model_data(self) -> bool:
"""Send model data to the detector."""
try:
model_path = self.detector_config.model.path
if not os.path.exists(model_path):
logger.error(f"Model file not found: {model_path}")
return False
logger.info(f"Transferring model to detector: {self._model_name}")
with open(model_path, "rb") as f:
model_data = f.read()
header = {"model_data": True, "model_name": self._model_name}
header_bytes = json.dumps(header).encode("utf-8")
self._socket.send_multipart([header_bytes, model_data])
# Temporarily increase timeout for model loading (can take several seconds)
original_timeout = self._socket.getsockopt(zmq.RCVTIMEO)
self._socket.setsockopt(zmq.RCVTIMEO, 30000)
try:
# Receive response
response_frames = self._socket.recv_multipart()
finally:
# Restore original timeout
self._socket.setsockopt(zmq.RCVTIMEO, original_timeout)
# Check if this is a JSON response (model management)
if len(response_frames) == 1:
try:
response = json.loads(response_frames[0].decode("utf-8"))
model_saved = response.get("model_saved", False)
model_loaded = response.get("model_loaded", False)
if model_saved and model_loaded:
logger.info(
f"Model {self._model_name} transferred and loaded successfully"
)
else:
logger.error(
f"Model transfer failed: saved={model_saved}, loaded={model_loaded}"
)
return model_saved and model_loaded
except json.JSONDecodeError:
logger.warning("Received non-JSON response for model data transfer")
return False
else:
logger.warning(
"Received unexpected response format for model data transfer"
)
return False
except Exception as e:
logger.error(f"Failed to send model data: {e}")
return False
def _read_offload_support(self, response: dict[str, Any]) -> None:
"""Record whether the detector applies the input transform itself."""
offload = response.get("input_transform_offload", False) is True
if offload != self.input_transform_offload:
self.input_transform_offload = offload
logger.info(
f"Detector input transform offload: {'on' if offload else 'off'}"
)
def _build_header(self, tensor_input: np.ndarray) -> bytes:
header: dict[str, Any] = {
"shape": list(tensor_input.shape),
"dtype": str(tensor_input.dtype.name),
"model_type": str(self.detector_config.model.model_type.name),
}
if self.input_transform_offload:
header["input_tensor"] = self.detector_config.model.input_tensor.value
header["input_dtype"] = self.detector_config.model.input_dtype.value
return json.dumps(header).encode("utf-8")
def _decode_response(self, frames: list[bytes]) -> np.ndarray:
try:
if len(frames) == 1:
# Single-frame raw float32 (20x6)
buf = frames[0]
if len(buf) != 20 * 6 * 4:
logger.warning(
f"ZMQ detector received unexpected payload size: {len(buf)}"
)
return self._zero_result
return np.frombuffer(buf, dtype=np.float32).reshape((20, 6))
if len(frames) >= 2:
header = json.loads(frames[0].decode("utf-8"))
shape = tuple(header.get("shape", []))
dtype = np.dtype(header.get("dtype", "float32"))
return np.frombuffer(frames[1], dtype=dtype).reshape(shape)
logger.warning("ZMQ detector received empty reply")
return self._zero_result
except Exception as exc: # noqa: BLE001
logger.error(f"ZMQ detector failed to decode response: {exc}")
return self._zero_result
def detect_raw(self, tensor_input: np.ndarray) -> np.ndarray:
if not self._model_ready:
logger.warning("Model not ready, returning zero detections")
return self._zero_result
try:
header_bytes = self._build_header(tensor_input)
payload_bytes = memoryview(tensor_input.tobytes(order="C"))
# Send request
self._socket.send_multipart([header_bytes, payload_bytes])
# Receive reply
reply_frames = self._socket.recv_multipart()
detections = self._decode_response(reply_frames)
# Ensure output shape and dtype are exactly as expected
return detections
except zmq.Again:
# Timeout
logger.debug("ZMQ detector request timed out; resetting socket")
try:
self._create_socket()
self._initialize_model()
except Exception:
pass
return self._zero_result
except zmq.ZMQError as exc:
logger.error(f"ZMQ detector ZMQError: {exc}; resetting socket")
try:
self._create_socket()
self._initialize_model()
except Exception:
pass
return self._zero_result
except Exception as exc: # noqa: BLE001
logger.error(f"ZMQ detector unexpected error: {exc}")
return self._zero_result
def __del__(self) -> None: # pragma: no cover - best-effort cleanup
try:
if self._socket is not None:
self._socket.close(linger=self.detector_config.linger_ms)
except Exception:
pass