Instructions to use camstack/camstack-trainer with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- ultralytics
How to use camstack/camstack-trainer with ultralytics:
from huggingface_hub import hf_hub_download from ultralytics import YOLO # pick the weights file from this repo's "Files and versions" tab weights = hf_hub_download("camstack/camstack-trainer", "<weights>.pt") model = YOLO(weights) source = 'http://images.cocodataset.org/val2017/000000039769.jpg' model.predict(source=source, save=True) - Notebooks
- Google Colab
- Kaggle
camstack-train 0.1.0: YOLOv9 root detector fine-tune (job.json in, result.json + bundle.tar out)
ebeed67 verified Download camstack_train/_vendor/pipeline_python/inference_pool.py from camstack/camstack-trainer: direct link, hf CLI and curl.
- Browser
- Download file 158 kB
-
https://huggingface.co/camstack/camstack-trainer/resolve/main/camstack_train/_vendor/pipeline_python/inference_pool.py
- Command line
-
hf download hf://camstack/camstack-trainer/camstack_train/_vendor/pipeline_python/inference_pool.py
-
curl -L -o inference_pool.py https://huggingface.co/camstack/camstack-trainer/resolve/main/camstack_train/_vendor/pipeline_python/inference_pool.py
158 kB
| #!/usr/bin/env python3 | |
| """Async inference pool — request_id multiplexed, per-runtime concurrency. | |
| Per-runtime ML backends (coreml / openvino / onnx): | |
| - asyncio main loop reads requests from stdin; inference is dispatched | |
| onto a runtime-specific executor so the reader never blocks. | |
| - Each request carries a 32-bit id; responses are tagged with the same | |
| id so the Node side can keep N requests in flight concurrently. | |
| - Runtime executors: | |
| CoreML → ThreadPoolExecutor(1) — ANE is single-context; one | |
| Python thread is enough, and avoids GIL thrashing. | |
| OpenVINO → ThreadPoolExecutor(OPTIMAL_NUMBER_OF_INFER_REQUESTS), | |
| each worker checking out an InferRequest from a fixed | |
| per-model pool (created once, post-compile), with the model | |
| compiled on ONE explicit device under the LATENCY hint (no | |
| AUTO — its CPU->GPU handover invalidated in-flight requests). | |
| ONNX → ThreadPoolExecutor(N) — N independent InferenceSessions | |
| where N = concurrency setting; each session pinned to | |
| its own worker thread. | |
| - Optional raw-frame path skips JPEG decode when the caller already | |
| has a decoded RGB/BGR/GRAY buffer (e.g. the stream-broker decoder | |
| cap output). | |
| Startup protocol (Node → Python, v2): | |
| 1. [4B total_len][4B req_id=0][1B msg_type=0x00][JSON config] | |
| config = {"runtime": ..., "models": [...], "concurrency": N?} | |
| 2. Python responds with [4B total_len][4B req_id=0][JSON ready status]. | |
| Runtime protocol: | |
| Request: [4B total_len][4B req_id][1B msg_type][payload] | |
| Response: [4B total_len][4B req_id][JSON payload] | |
| msg_type: | |
| 0x00 — command payload = JSON | |
| 0x01 — infer_jpeg payload = [1B model_idx][JPEG bytes] | |
| 0x02 — infer_raw payload = [1B model_idx][4B width][4B height] | |
| [1B fmt 0=RGB,1=BGR,2=GRAY][pixels] | |
| Wire-protocol versioning (D350): the startup config may declare | |
| `protocolVersion`; the effective wire version is `min(PROTOCOL_VERSION, | |
| declared)` and the ready reply reports this worker's own PROTOCOL_VERSION so | |
| the Node side can negotiate the same minimum. At v2, every SHEDDABLE request | |
| (infer_jpeg / infer_raw / infer_batch / infer_cached) prefixes its payload | |
| with the request's ABSOLUTE deadline — [8B epoch-ms uint64 LE] — and this | |
| worker checks it immediately before predict, answering | |
| `{"dropped": true, "shedReason": "deadline-expired"}` instead of executing | |
| work whose caller has already abandoned it (the 2026-09-04 storm: 4h15m of | |
| the accelerator at 100% on replies discarded as `unknown request id`). | |
| Commands and model loads NEVER carry a deadline at any version — a shed | |
| command desynchronizes model state. A host that declares nothing is an OLD | |
| host and gets byte-identical v1 behaviour; skew in either direction is the | |
| NORMAL state during a rolling deploy. | |
| Commands: load, unload, replace, reconfigure, status, uncache_frame, mem_stats. | |
| `load`/`replace` compile on their own thread under a hard bound and never block | |
| the loop; a load past the bound answers `reason: "compile-timeout"` and poisons | |
| the worker (D653, see `ModelLoader`). | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import concurrent.futures | |
| import contextlib | |
| import fcntl | |
| import gc | |
| import io | |
| import json | |
| import os | |
| import platform | |
| import re | |
| import shutil | |
| import struct | |
| import sys | |
| import threading | |
| import time | |
| from collections import deque | |
| from dataclasses import dataclass, field | |
| from typing import Any, Awaitable, Callable, Optional | |
| import numpy as np | |
| from PIL import Image | |
| from postprocessors import POSTPROCESSORS | |
| # --------------------------------------------------------------------------- | |
| # Wire protocol constants | |
| # --------------------------------------------------------------------------- | |
| MSG_COMMAND = 0x00 | |
| MSG_INFER_JPEG = 0x01 | |
| MSG_INFER_RAW = 0x02 | |
| # MSG_INFER_BATCH — N raw items packed into a single IPC frame. | |
| # Wire format: | |
| # [1B model_idx][1B count][N × item] | |
| # Each item: | |
| # [4B width][4B height][1B fmt][4B size][size bytes raw] | |
| # Response payload: | |
| # {"results": [<single-detect dict>, ...]} (same length as count) | |
| # Each item is dispatched concurrently via asyncio.gather so the | |
| # predict pool's existing parallelism applies; the saving over N | |
| # separate calls is one IPC round-trip per item collapsed to one, | |
| # with batch=4 semantics for fair benchmarking. | |
| MSG_INFER_BATCH = 0x03 | |
| MSG_CACHE_FRAME = 0x04 | |
| MSG_INFER_CACHED = 0x05 | |
| RAW_FMT_RGB = 0x00 | |
| RAW_FMT_BGR = 0x01 | |
| RAW_FMT_GRAY = 0x02 | |
| # --------------------------------------------------------------------------- | |
| # Wire-protocol versioning + the deadline on the wire (D350) | |
| # --------------------------------------------------------------------------- | |
| # Highest wire version THIS worker implements. Reported in the ready reply; | |
| # the effective version per connection is min() of both sides' declarations, | |
| # so a worker and a host at different versions (normal during a rolling | |
| # deploy) always agree on the layout. | |
| PROTOCOL_VERSION = 2 | |
| # Opcodes that may carry (and be shed by) a deadline. Mirrors | |
| # SHEDDABLE_MSG_TYPES in shared-inference-pool.ts — commands and cacheFrame | |
| # are absent on BOTH sides: a shed command desynchronizes model state, a shed | |
| # cacheFrame strands the inferCached that follows it. | |
| SHEDDABLE_MSG_TYPES = frozenset({MSG_INFER_JPEG, MSG_INFER_RAW, MSG_INFER_BATCH, MSG_INFER_CACHED}) | |
| # Bytes of the absolute-deadline prefix a v2 sheddable payload carries. | |
| DEADLINE_PREFIX_LEN = 8 | |
| # Sampling stride for the deadline-expired shed stderr line. Under a storm | |
| # this fires at frame rate; one line per shed buried every other signal the | |
| # last time a shed path logged unsampled (45 741 rows/hour, 2026-08-01). The | |
| # counter (PoolCounters.shed_deadline_expired) keeps the exact total. | |
| DEADLINE_SHED_LOG_EVERY = 500 | |
| def _negotiate_wire_version(config: dict) -> int: | |
| """Effective wire version for this connection: min(ours, host declared). | |
| No declaration (an OLD host) or a garbage value negotiates v1 — the | |
| byte-identical historical layout. Pure; unit-tested. | |
| """ | |
| try: | |
| declared = int(config.get("protocolVersion", 1) or 1) | |
| except (TypeError, ValueError): | |
| declared = 1 | |
| return max(1, min(PROTOCOL_VERSION, declared)) | |
| def _split_deadline(msg_type: int, payload: bytes, wire_version: int) -> "tuple[int, bytes]": | |
| """Split a request payload into (deadline_ms, rest) by the NEGOTIATED | |
| version. v1, and every non-sheddable opcode at any version, passes the | |
| payload through untouched with deadline 0 (= no deadline, never sheds). | |
| Raises ValueError on a v2 sheddable payload too short to carry its | |
| prefix. Pure; unit-tested. | |
| """ | |
| if wire_version < 2 or msg_type not in SHEDDABLE_MSG_TYPES: | |
| return 0, payload | |
| if len(payload) < DEADLINE_PREFIX_LEN: | |
| raise ValueError( | |
| f"truncated deadline prefix on msg_type={msg_type} " | |
| f"({len(payload)} bytes, need {DEADLINE_PREFIX_LEN})" | |
| ) | |
| (deadline_ms,) = struct.unpack("<Q", payload[:DEADLINE_PREFIX_LEN]) | |
| return deadline_ms, payload[DEADLINE_PREFIX_LEN:] | |
| def _deadline_expired(deadline_ms: int, now_ms: float) -> bool: | |
| """0 means "no deadline" and never expires (v1 host / unstamped).""" | |
| return deadline_ms > 0 and now_ms > deadline_ms | |
| def _dropped_deadline_payload() -> dict: | |
| """The shed reply for an expired request — same shape the TS side already | |
| recognises as flow control (`poolResultToEngineOutput`).""" | |
| return { | |
| "kind": "detections", | |
| "detections": [], | |
| "inferenceMs": 0, | |
| "dropped": True, | |
| "shedReason": "deadline-expired", | |
| } | |
| def _log_deadline_shed(total: int, model_idx: int, late_ms: float) -> None: | |
| """Sampled stderr line for a deadline shed — first occurrence and every | |
| DEADLINE_SHED_LOG_EVERY-th, carrying the running total so nothing is | |
| lost, only repeated. A branch that drops work silently is how the | |
| 2026-09-04 collapse produced not one line in four hours.""" | |
| if total != 1 and total % DEADLINE_SHED_LOG_EVERY != 0: | |
| return | |
| sys.stderr.write( | |
| f"deadline-expired shed: model {model_idx}, {late_ms:.0f}ms past deadline " | |
| f"(total {total}, sampled every {DEADLINE_SHED_LOG_EVERY})\n" | |
| ) | |
| sys.stderr.flush() | |
| # --------------------------------------------------------------------------- | |
| # Preprocessing (unchanged from v1 — same shapes/math, moved into helpers) | |
| # --------------------------------------------------------------------------- | |
| def letterbox_image(img: Image.Image, size: int) -> tuple[Image.Image, float, tuple[int, int]]: | |
| """Resize with letterbox padding; returns (PIL canvas, scale, (padX, padY)). | |
| Stays in uint8 PIL space — no float conversion. Used by the imageType | |
| fast path where CoreML's predict() accepts a PIL Image directly and | |
| handles the BGR/normalize/CVPixelBuffer dance internally. | |
| """ | |
| w, h = img.size | |
| scale = size / max(w, h) | |
| new_w = int(w * scale) | |
| new_h = int(h * scale) | |
| pad_x = (size - new_w) // 2 | |
| pad_y = (size - new_h) // 2 | |
| if w == size and h == size: | |
| # Already at target size — skip resize + paste + alloc entirely. | |
| return img, 1.0, (0, 0) | |
| # A frameRef view already arrives at the model's content dimensions | |
| # (e.g. 320x180 for a 16:9 YOLO-320 input). Add only the square padding; | |
| # repeating the same-size geometric resize wastes CPU and can blur pixels. | |
| resized = img if (w, h) == (new_w, new_h) else img.resize((new_w, new_h), Image.BILINEAR) | |
| canvas = Image.new("RGB", (size, size), (114, 114, 114)) | |
| canvas.paste(resized, (pad_x, pad_y)) | |
| return canvas, scale, (pad_x, pad_y) | |
| def letterbox(img: Image.Image, size: int) -> tuple[np.ndarray, float, tuple[int, int]]: | |
| """Resize with letterbox padding; returns (array[0..1] HWC float32, scale, (padX, padY)). | |
| Used by multiArrayType inputs where the model expects a normalized | |
| float tensor. Does the uint8→float32/255 conversion exactly once. | |
| """ | |
| canvas, scale, pad = letterbox_image(img, size) | |
| arr = np.array(canvas, dtype=np.float32) / 255.0 | |
| return arr, scale, pad | |
| def resize_image(img: Image.Image, width: int, height: int) -> np.ndarray: | |
| resized = img.resize((width, height), Image.BILINEAR) | |
| return np.array(resized, dtype=np.float32) / 255.0 | |
| def decode_jpeg(jpeg: bytes) -> Image.Image: | |
| return Image.open(io.BytesIO(jpeg)).convert("RGB") | |
| def wrap_raw(raw: bytes, width: int, height: int, fmt: int) -> Image.Image: | |
| """Zero-copy (via PIL frombuffer) wrap of a raw frame buffer.""" | |
| if fmt == RAW_FMT_GRAY: | |
| img = Image.frombuffer("L", (width, height), raw, "raw", "L", 0, 1) | |
| return img.convert("RGB") | |
| mode = "BGR" if fmt == RAW_FMT_BGR else "RGB" | |
| img = Image.frombuffer("RGB", (width, height), raw, "raw", mode, 0, 1) | |
| return img | |
| # --------------------------------------------------------------------------- | |
| # Model slot + runtime-specific loading | |
| # --------------------------------------------------------------------------- | |
| class ModelSlot: | |
| model: Any = None | |
| config: dict = field(default_factory=dict) | |
| loaded: bool = False | |
| predict_fn: Optional[Callable[[dict], dict]] = None | |
| # Atomic snapshot of model.spec input names; used by both single | |
| # and batched predict paths to filter input_dicts against the | |
| # CURRENTLY-loaded model. Defends against the race where an in- | |
| # flight preprocess holds a reference to the previously-loaded | |
| # model's `_input_names` and adds keys (iouThreshold, | |
| # confidenceThreshold) that the new model rejects. | |
| input_names: frozenset = field(default_factory=frozenset) | |
| _runtime: str = "" | |
| _runtime_lib: Any = None | |
| # Max OPTIMAL_NUMBER_OF_INFER_REQUESTS across loaded OpenVINO models — used | |
| # to size the predict pool so the per-model InferRequest pool is actually fed. | |
| _ov_optimal_reqs: int = 0 | |
| # Default OpenVINO predict-pool size when models load lazily (so we can't yet | |
| # query the device's optimal request count). ≈ measured optimal on Intel | |
| # CPU/iGPU/NPU (4-5). Threads idle-block on infer, so over-provisioning is cheap. | |
| OV_DEFAULT_CONCURRENCY: int = 4 | |
| # Floor for the InferRequest / predict-pool size when the device DOES answer. | |
| # Under PERFORMANCE_HINT=LATENCY an Intel iGPU typically reports 1-2: one | |
| # execution stream, and the extra requests buy pipelining the LATENCY hint is | |
| # not asking for. Two rather than one so a second frame can be preprocessed | |
| # while the first infers, without paying for four. | |
| OV_MIN_CONCURRENCY: int = 2 | |
| def _ov_request_count(optimal: int) -> int: | |
| """How many InferRequests to create up front. | |
| ONE function because two call sites must agree: the per-model request pool | |
| in `_load_model` and the predict-pool worker count in `_run`. They used to | |
| compute it separately with the same expression, and the expression was | |
| `max(optimal, OV_DEFAULT_CONCURRENCY)` — a FLOOR of four, not the fallback | |
| its own comment described. On this hub that meant every model held four | |
| GPU-resident InferRequests forever, per pool, while the device was asking | |
| for one or two and `intel_gpu_top` showed the GPU 92% idle. Creating a | |
| VAAPI device for an HEVC decode then failed with ENOMEM. | |
| `optimal <= 0` means the device could not be asked — models load lazily, | |
| the property is missing — and UNKNOWN is not SMALL: that case keeps the | |
| old default rather than guessing low. | |
| """ | |
| if optimal <= 0: | |
| return OV_DEFAULT_CONCURRENCY | |
| return max(optimal, OV_MIN_CONCURRENCY) | |
| # OpenVINO device decision. AUTO is NEVER used: the AUTO plugin runs inference | |
| # on a CPU "helper" while the GPU JIT-compiles, then hands the model over to the | |
| # GPU mid-flight. InferRequests created during that window bind to the | |
| # CPU-helper's Result ports; after handover the served model's ports diverge and | |
| # req.infer() raises "Cannot find tensor for port opset1::Result score_8" (proven | |
| # live on the N100 iGPU). Compiling ONE explicit device removes the handover | |
| # entirely. We resolve an ordered candidate list (NPU > GPU > CPU) and compile the | |
| # first that succeeds, with CPU as the guaranteed floor. | |
| # | |
| # OpenVINO PERFORMANCE_HINT is LATENCY: one deterministic execution stream. The | |
| # THROUGHPUT hint multiplied memory and lengthened tail latency (an N100 | |
| # OOM-freeze factor, see the ov_config note in _load_model); its env switch and | |
| # the never-enabled AsyncInferQueue path were removed on 2026-10-06. | |
| _OV_PERF_HINT: str = "LATENCY" | |
| # OpenVINO COMPILATION thread pool (CAMSTACK_OV_COMPILATION_THREADS, default 4). | |
| # | |
| # OpenVINO sizes the pool that COMPILES a model to the machine's core count, per | |
| # compile, and TBB never retires those workers — they live as long as the | |
| # process. A container cgroup counts THREADS, not processes, so every compile | |
| # permanently spends `nproc` of the cgroup's pid budget. | |
| # | |
| # Measured on the hub (20 cores) 2026-09-18, READ-ONLY, before this bound: | |
| # pids.current 991 / pids.max 2048 (22 fork refusals in 45 minutes earlier). | |
| # The detection pool (pid 1183) held 214 threads 17 minutes after boot — the | |
| # single largest thread holder in the container, 3x the next one (ffmpeg, 80). | |
| # Thread start times cluster EXACTLY on its compile timestamps: 167 threads in | |
| # a 4s boot burst (~8 models), then 47 more in a 2-model burst ten minutes | |
| # later. 23 threads per compile on a 20-core box. ~10 models are resident. | |
| # Expected after: ~4 threads per compile, i.e. ~40 instead of ~214 in that pool. | |
| # | |
| # Four, not one: the cold-compile storm (~25s for the full model set on an N100 | |
| # iGPU) is where frames pile up, so the bound keeps some parallelism. It costs | |
| # nothing at inference time — COMPILATION_NUM_THREADS is not | |
| # INFERENCE_NUM_THREADS, which is deliberately left at OpenVINO's default (see | |
| # `_ov_compilation_properties`). 0 disables the bound entirely. | |
| try: | |
| OV_COMPILATION_NUM_THREADS: int = int( | |
| os.environ.get("CAMSTACK_OV_COMPILATION_THREADS", "4") | |
| ) | |
| except ValueError: | |
| OV_COMPILATION_NUM_THREADS = 4 | |
| def _ov_supported_properties(core: Any, device: str) -> list[str]: | |
| """SUPPORTED_PROPERTIES for `device`, or [] when the device cannot be asked. | |
| A probe that FAILS knows nothing, and unknown is not "supported": the caller | |
| then sends no extra property and compiles exactly as it always did. | |
| """ | |
| try: | |
| return [str(p) for p in core.get_property(device, "SUPPORTED_PROPERTIES")] | |
| except Exception: | |
| return [] | |
| def _ov_compilation_properties(supported: Any, threads: int) -> dict: | |
| """The compile-time properties that BOUND OpenVINO's compilation pool. | |
| Only `COMPILATION_NUM_THREADS`, and only for a device that lists it. A | |
| plugin REJECTS a property it does not support and the rejection surfaces as | |
| a compile exception — which the failover loop in `_load_model` would read as | |
| "this device cannot run the model" and quietly land a GPU model on the CPU. | |
| Absent beats rejected. | |
| INFERENCE_NUM_THREADS is deliberately NOT set here. Under PERFORMANCE_HINT | |
| =LATENCY the compiled model runs one stream, and the CPU device is this | |
| pool's guaranteed fallback floor — capping its inference threads would trade | |
| a pid saving we have not measured for a latency regression on exactly the | |
| nodes that have no accelerator. The compilation pool is pure overhead; the | |
| inference pool is the work. | |
| """ | |
| if threads <= 0 or "COMPILATION_NUM_THREADS" not in set(supported): | |
| return {} | |
| return {"COMPILATION_NUM_THREADS": threads} | |
| # Inference precisions a catalog entry may require on an OpenVINO GPU. | |
| _OV_GPU_PRECISIONS = ("f32", "f16") | |
| def _ov_precision_properties(config: dict, device: str, supported: Any) -> dict: | |
| """The INFERENCE_PRECISION_HINT a model's catalog entry requires on a GPU. | |
| An Intel GPU runs an FP32 IR at f16 by default. Some IRs do not survive | |
| that: on 2026-10-07 the YOLO26 base IRs `yolo26s`, `yolo26m` and | |
| `yolo26m-320` returned no valid detection on the Core Ultra iGPU and the | |
| N100 (0 true / 146 false positives, or NaN scores) and matched the CPU at | |
| f32. The entry says so (`gpuInferencePrecision`, threaded as config key | |
| `gpuInferencePrecision`) — never a model-name match. | |
| GPU only: the NPU and CPU already compute these IRs correctly. And only for | |
| a device that lists the key — a rejected property would make the failover | |
| loop push the model onto the next device. | |
| """ | |
| precision = config.get("gpuInferencePrecision") | |
| if precision not in _OV_GPU_PRECISIONS: | |
| return {} | |
| if not str(device).upper().startswith("GPU"): | |
| return {} | |
| if "INFERENCE_PRECISION_HINT" not in set(supported): | |
| return {} | |
| return {"INFERENCE_PRECISION_HINT": precision} | |
| def _resolve_ov_devices(available: list[str], full_names: dict[str, str]) -> list[str]: | |
| """Ordered EXPLICIT device candidates for compile-time failover (NPU>GPU>CPU). | |
| Never returns an AUTO/AUTO:.. mode (see the module comment above). `available` | |
| is `core.available_devices`; `full_names` maps each device to its | |
| FULL_DEVICE_NAME (missing entries = the probe failed, left unclassified). A | |
| GPU is only accepted when its FULL_DEVICE_NAME probed successfully (this is | |
| also how an Intel iGPU vs an NVIDIA dGPU is told apart downstream). Ordering: | |
| NPU(s) -> GPU(s) -> CPU (always last). The compile loop tries each in order | |
| and stops at the first that compiles; CPU is the guaranteed fallback. | |
| """ | |
| npus: list[str] = [] | |
| gpus: list[str] = [] | |
| for device in available: | |
| full_device_name = full_names.get(device) | |
| if "NPU" in device: | |
| npus.append(device) | |
| elif "GPU" in device and full_device_name is not None: | |
| gpus.append(device) | |
| ordered: list[str] = [*npus, *gpus] | |
| if "CPU" in available and "CPU" not in ordered: | |
| ordered.append("CPU") | |
| if not ordered: | |
| ordered = ["CPU"] | |
| return ordered | |
| # ONNX Runtime execution-provider names used by _resolve_onnx_providers. | |
| ONNX_CUDA_EP = "CUDAExecutionProvider" | |
| ONNX_COREML_EP = "CoreMLExecutionProvider" | |
| ONNX_CPU_EP = "CPUExecutionProvider" | |
| def _resolve_onnx_providers( | |
| device: str, | |
| plat: str, | |
| machine: str, | |
| available: list[str], | |
| cuda_device_id: int = 0, | |
| ) -> list: | |
| """Ordered ONNX Runtime providers — hardware first, CPU always last. | |
| Ordering for the AUTO/default case: | |
| darwin -> CoreMLExecutionProvider | |
| linux/win on x86_64/AMD64 -> ("CUDAExecutionProvider", {"device_id": N}) | |
| always -> CPUExecutionProvider appended LAST | |
| Deliberate adaptations for camstack: | |
| - an operator device pin is honored: "cuda"/"coreml" put that EP first | |
| (CPU still last as fallback); "cpu" compiles CPU only. | |
| - hardware entries not present in `available` | |
| (ort.get_available_providers()) are dropped — camstack ships the | |
| plain `onnxruntime` wheel on most nodes, and requesting an EP the | |
| build lacks raises at session creation instead of falling back. | |
| - platform match uses startswith so a substring test on "win" never | |
| wrongly matches "darwin". | |
| """ | |
| dev = (device or "").strip().lower() | |
| ordered: list = [] | |
| if dev == "cuda": | |
| ordered.append((ONNX_CUDA_EP, {"device_id": cuda_device_id})) | |
| elif dev == "coreml": | |
| ordered.append(ONNX_COREML_EP) | |
| elif dev == "cpu": | |
| pass # explicit CPU pin — no hardware EPs | |
| else: | |
| # AUTO/default — platform-driven ordering. | |
| if plat == "darwin": | |
| ordered.append(ONNX_COREML_EP) | |
| if plat.startswith(("linux", "win")) and machine in ("x86_64", "AMD64"): | |
| ordered.append((ONNX_CUDA_EP, {"device_id": cuda_device_id})) | |
| providers: list = [ | |
| p for p in ordered | |
| if (p[0] if isinstance(p, tuple) else p) in available | |
| ] | |
| providers.append(ONNX_CPU_EP) | |
| return providers | |
| # --------------------------------------------------------------------------- | |
| # CoreML persistent compiled-model cache | |
| # --------------------------------------------------------------------------- | |
| # Known macOS behavior: the OS-level compiled-model cache is NOT reused across | |
| # process restarts — every | |
| # `ct.models.MLModel(.mlpackage)` load RECOMPILES the model, and the stale OS | |
| # cache is only cleared on reboot. camstack loads the .mlpackage on every pool | |
| # spawn, so each spawn paid the full recompile. Fix: after the first compile, | |
| # copy the compiled `.mlmodelc` (get_compiled_model_path() returns a TEMP dir | |
| # that dies with the MLModel object) into a persistent `.coreml-cache` dir next | |
| # to the model (mirrors the OpenVINO `.ov-cache` convention); subsequent spawns | |
| # load it via `ct.models.CompiledMLModel(...)` — no recompile. A stamp file | |
| # fingerprinting the source model invalidates the cache when the model changes. | |
| # HARD RULE: every failure in this machinery falls back to the plain | |
| # `ct.models.MLModel(path)` load — exactly the pre-cache behavior. | |
| # --------------------------------------------------------------------------- | |
| # Compile-cache lock — one worker compiles a model, the rest wait and HIT | |
| # --------------------------------------------------------------------------- | |
| # | |
| # A pool is several python PROCESSES. They start together, all miss the | |
| # compiled-model cache because nobody has written it yet, all pay the compile, | |
| # and then all try to replace the same directory. Measured on the operator's | |
| # Mac on 2026-09-19 over three hours: **59 cache HITs, 55 MISSes, 5 persist | |
| # failures** — half the loads recompiling a model that was already on disk. | |
| # | |
| # The failures name a path INSIDE the bundle: | |
| # | |
| # CoreML: could not persist compiled-model cache | |
| # ([Errno 2] No such file or directory: '.../yolov9m-320.mlmodelc/weights') | |
| # | |
| # which is `_persist_coreml_compiled` removing the old cache while another | |
| # worker swaps its own copy into the same place: `rmtree` walks a directory and | |
| # a child ceases to exist under it. | |
| # | |
| # So this lock is not belt-and-braces. It is what turns N compiles of one model | |
| # into one, and what makes "replace the cache" single-writer. | |
| # | |
| # It is deliberately NOT fatal in any direction: an unusable lock directory, a | |
| # platform without `flock`, or a holder that died mid-compile all end the same | |
| # way — the caller compiles without the cache, which is exactly what it did | |
| # before the cache existed. A cache is an optimisation, and an optimisation | |
| # that can refuse to serve a camera is worse than every miss it prevents. | |
| # How long a waiter gives the compiling worker before compiling itself. | |
| COMPILE_CACHE_LOCK_TIMEOUT_SEC = 180.0 | |
| _COMPILE_CACHE_LOCK_POLL_SEC = 0.05 | |
| def compile_cache_lock( | |
| cache_root: str, | |
| stem: str, | |
| timeout_sec: float = COMPILE_CACHE_LOCK_TIMEOUT_SEC, | |
| ): | |
| """Hold the per-MODEL compile lock, or yield False having failed to. | |
| Per model, not per cache: compiling the detector must not serialise behind | |
| the embedder, which is a different file and a different compile. | |
| """ | |
| handle = None | |
| held = False | |
| try: | |
| os.makedirs(cache_root, exist_ok=True) | |
| handle = open(os.path.join(cache_root, f"{stem}.lock"), "a+") | |
| deadline = time.monotonic() + max(0.0, timeout_sec) | |
| while True: | |
| try: | |
| fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) | |
| held = True | |
| break | |
| except OSError: | |
| if time.monotonic() >= deadline: | |
| break | |
| time.sleep(_COMPILE_CACHE_LOCK_POLL_SEC) | |
| except (OSError, AttributeError, NameError): | |
| held = False | |
| try: | |
| yield held | |
| finally: | |
| if handle is not None: | |
| try: | |
| if held: | |
| fcntl.flock(handle.fileno(), fcntl.LOCK_UN) | |
| except OSError: | |
| pass | |
| try: | |
| handle.close() | |
| except OSError: | |
| pass | |
| COREML_CACHE_DIR_NAME = ".coreml-cache" | |
| COREML_CACHE_STAMP_VERSION = 1 | |
| def _coreml_cache_paths(model_path: str) -> tuple[str, str]: | |
| """(compiled .mlmodelc cache dir, stamp file) for a source model path. | |
| The cache lives in a dot-dir next to the model file — same convention as | |
| the OpenVINO `.ov-cache` — so it ships/wipes with the models dir. | |
| """ | |
| models_dir = os.path.dirname(os.path.abspath(model_path)) | |
| stem = os.path.splitext(os.path.basename(model_path))[0] | |
| cache_root = os.path.join(models_dir, COREML_CACHE_DIR_NAME) | |
| return ( | |
| os.path.join(cache_root, stem + ".mlmodelc"), | |
| os.path.join(cache_root, stem + ".stamp.json"), | |
| ) | |
| def _coreml_model_fingerprint(model_path: str) -> dict: | |
| """Cheap source-model fingerprint — file count + total size + max mtime. | |
| Works for both a plain file (.mlmodel) and a .mlpackage DIRECTORY (the | |
| normal case), whose top-level mtime alone would not reflect content edits. | |
| JSON-serializable so it round-trips through the stamp file unchanged. | |
| The package's top-level `Manifest.json` counts by size only, never by | |
| mtime: coremltools rewrites it (same bytes) on EVERY open of the package — | |
| `MLModel` and `utils.load_spec` alike (measured, coremltools 8.3, macOS | |
| 26.2). The cache HIT path calls `load_spec`, so keyed on that mtime every | |
| HIT invalidated the stamp for the next spawn: HIT, MISS, HIT, MISS. A | |
| re-download rewrites the model and weights files too, so it still shows. | |
| """ | |
| if not os.path.isdir(model_path): | |
| st = os.stat(model_path) | |
| return {"kind": "file", "files": 1, "size": st.st_size, "mtimeNs": st.st_mtime_ns} | |
| total_size = 0 | |
| latest_mtime = 0 | |
| file_count = 0 | |
| manifest = os.path.join(model_path, "Manifest.json") | |
| for root, _dirs, files in os.walk(model_path): | |
| for name in files: | |
| file_path = os.path.join(root, name) | |
| st = os.stat(file_path) | |
| total_size += st.st_size | |
| file_count += 1 | |
| if file_path == manifest: | |
| continue | |
| if st.st_mtime_ns > latest_mtime: | |
| latest_mtime = st.st_mtime_ns | |
| return {"kind": "dir", "files": file_count, "size": total_size, "mtimeNs": latest_mtime} | |
| def _coreml_stamp_matches(stamp: Any, fingerprint: dict) -> bool: | |
| """Pure decision: does a parsed stamp validate the cached compile?""" | |
| return ( | |
| isinstance(stamp, dict) | |
| and stamp.get("version") == COREML_CACHE_STAMP_VERSION | |
| and stamp.get("fingerprint") == fingerprint | |
| ) | |
| def _persist_coreml_compiled( | |
| compiled_src: str, cache_path: str, stamp_path: str, fingerprint: dict, | |
| ) -> None: | |
| """Copy coremltools' TEMP compiled dir into the persistent cache path. | |
| Crash-safe ordering: copy into a tmp sibling on the SAME filesystem, then | |
| `os.replace` into place, and write the stamp LAST — a stamp file only ever | |
| exists for a fully-moved `.mlmodelc`, so a crash mid-copy can never | |
| produce a half-written cache that validates. Raises on failure — the | |
| caller logs and continues WITHOUT a cache (model already loaded fine). | |
| """ | |
| cache_root = os.path.dirname(cache_path) | |
| os.makedirs(cache_root, exist_ok=True) | |
| tmp_dir = f"{cache_path}.tmp-{os.getpid()}" | |
| if os.path.exists(tmp_dir): | |
| shutil.rmtree(tmp_dir) | |
| try: | |
| shutil.copytree(compiled_src, tmp_dir) | |
| # Drop any stale previous cache: its stamp first (so nothing validates | |
| # the old dir), then the dir itself — POSIX rename cannot replace a | |
| # non-empty directory. | |
| # Both removals tolerate "already gone". `compile_cache_lock` makes a | |
| # concurrent swap rare, not impossible — a worker killed between taking | |
| # the lock and the replace leaves precisely this state — and a previous | |
| # cache that has ALREADY disappeared is the outcome we wanted anyway. | |
| # Not tolerating it is what produced five | |
| # `[Errno 2] ... .mlmodelc/weights` failures in three hours. | |
| try: | |
| os.remove(stamp_path) | |
| except FileNotFoundError: | |
| pass | |
| shutil.rmtree(cache_path, ignore_errors=True) | |
| os.replace(tmp_dir, cache_path) | |
| except BaseException: | |
| shutil.rmtree(tmp_dir, ignore_errors=True) | |
| raise | |
| stamp_tmp = f"{stamp_path}.tmp-{os.getpid()}" | |
| try: | |
| with open(stamp_tmp, "w", encoding="utf-8") as fh: | |
| json.dump({"version": COREML_CACHE_STAMP_VERSION, "fingerprint": fingerprint}, fh) | |
| os.replace(stamp_tmp, stamp_path) | |
| except BaseException: | |
| try: | |
| os.remove(stamp_tmp) | |
| except OSError: | |
| pass | |
| raise | |
| def _try_coreml_cache_hit( | |
| ct: Any, path: str, compute_units: Any, | |
| ) -> tuple[Any, Any] | None: | |
| """(model, spec) from the persistent cache, or None for "compile it". | |
| Separate from the compile because it is asked TWICE: once before queueing | |
| for the compile lock, and again once the lock is held — the second ask is | |
| what turns every waiter's miss into a hit. | |
| `CompiledMLModel` has no `get_spec()`, so the spec is read straight from | |
| the source .mlpackage with `ct.utils.load_spec` (cheap — parses the | |
| protobuf, does not compile). Any failure answers None, which is the | |
| pre-cache behaviour exactly. | |
| """ | |
| try: | |
| cache_path, stamp_path = _coreml_cache_paths(path) | |
| fingerprint = _coreml_model_fingerprint(path) | |
| stamp: Any = None | |
| if os.path.isfile(stamp_path): | |
| with open(stamp_path, "r", encoding="utf-8") as fh: | |
| stamp = json.load(fh) | |
| if _coreml_stamp_matches(stamp, fingerprint) and os.path.isdir(cache_path): | |
| spec = ct.utils.load_spec(path) | |
| model = ct.models.CompiledMLModel(cache_path, compute_units) | |
| sys.stderr.write(f"CoreML: compiled-model cache HIT {cache_path}\n") | |
| sys.stderr.flush() | |
| return model, spec | |
| except Exception as exc: | |
| sys.stderr.write( | |
| f"CoreML: compiled-model cache lookup failed ({exc}); " | |
| f"falling back to plain compile of {os.path.basename(path)}\n" | |
| ) | |
| sys.stderr.flush() | |
| return None | |
| def _acquire_coreml_model(ct: Any, path: str, compute_units: Any) -> tuple[Any, Any]: | |
| """Return (model, spec) for the CoreML branch via the persistent cache. | |
| Cache HIT → `CompiledMLModel(cache)`, no recompile (`_try_coreml_cache_hit`). | |
| Cache MISS → take the per-model compile lock, ask AGAIN (the winner may | |
| have finished while we queued), then `MLModel(path)` — the | |
| pre-cache behaviour — and persist it for the next spawn. | |
| LOCK LOST → compile anyway, serve the model, write nothing. | |
| ANY failure in the cache machinery → the plain MLModel path, exactly as | |
| before this cache existed. Never raises differently from pre-cache code. | |
| """ | |
| hit = _try_coreml_cache_hit(ct, path, compute_units) | |
| if hit is not None: | |
| return hit | |
| # A MISS here is not yet a compile. The pool starts its workers together, | |
| # so every one of them arrives at this line at the same instant with the | |
| # same answer — measured 55 misses against 59 hits on the Mac, which is | |
| # not a cold cache, it is N workers each compiling what the others are | |
| # compiling. One of them takes the lock; the rest wait and then find the | |
| # HIT the winner just wrote. | |
| cache_root = os.path.dirname(_coreml_cache_paths(path)[0]) | |
| stem = os.path.splitext(os.path.basename(path))[0] | |
| with compile_cache_lock(cache_root, stem) as held: | |
| if held: | |
| hit = _try_coreml_cache_hit(ct, path, compute_units) | |
| if hit is not None: | |
| return hit | |
| model = ct.models.MLModel(path, compute_units=compute_units) | |
| spec = model.get_spec() | |
| # Only the lock HOLDER writes the cache. A worker that compiled | |
| # because it could not get the lock has a perfectly good model and no | |
| # business replacing a directory somebody else owns — that swap is the | |
| # race, and skipping it costs nothing: the holder is writing the same | |
| # bytes. | |
| if not held: | |
| sys.stderr.write( | |
| f"CoreML: compiled without the cache lock — compiled " | |
| f"{os.path.basename(path)} and left the cache to its holder\n" | |
| ) | |
| sys.stderr.flush() | |
| return model, spec | |
| try: | |
| cache_path, stamp_path = _coreml_cache_paths(path) | |
| fingerprint = _coreml_model_fingerprint(path) | |
| compiled_src = model.get_compiled_model_path() | |
| _persist_coreml_compiled(compiled_src, cache_path, stamp_path, fingerprint) | |
| sys.stderr.write( | |
| f"CoreML: compiled-model cache MISS — compiled and persisted {cache_path}\n" | |
| ) | |
| sys.stderr.flush() | |
| except Exception as exc: | |
| sys.stderr.write( | |
| f"CoreML: could not persist compiled-model cache ({exc}); " | |
| f"continuing without cache\n" | |
| ) | |
| sys.stderr.flush() | |
| return model, spec | |
| def _init_runtime(runtime: str, device: str = "") -> None: | |
| global _runtime, _runtime_lib | |
| _runtime = runtime | |
| if runtime == "coreml": | |
| import coremltools as ct | |
| _runtime_lib = ct | |
| elif runtime == "openvino": | |
| # `openvino.runtime` is deprecated (removed in a future OpenVINO release); | |
| # `Core` is exported from the top-level `openvino` namespace. | |
| from openvino import Core | |
| _runtime_lib = Core() | |
| elif runtime == "onnxruntime": | |
| import onnxruntime as ort | |
| _runtime_lib = ort | |
| elif runtime == "edgetpu": | |
| # Coral USB Edge TPU via LiteRT (ai-edge-litert) + the libedgetpu delegate. | |
| # The delegate is created ONCE here and shared across all model slots (the | |
| # USB device is a single hardware context). See docs: the delegate's TF | |
| # version MUST match the interpreter — feranick libedgetpu tf2.17.1 pairs | |
| # with ai-edge-litert 2.1.6. Install: /data/deps/edgetpu-install.sh. | |
| import threading | |
| from ai_edge_litert.interpreter import Interpreter, load_delegate | |
| lib_path = os.environ.get( | |
| "CAMSTACK_EDGETPU_LIB", "/data/deps/edgetpu/lib/libedgetpu.so.1", | |
| ) | |
| # Multiple Corals coexist and run concurrently, one pool each. libedgetpu | |
| # addresses a specific unit by an ordinal device string ("usb", "usb:1", | |
| # "pci:0", …); the orchestrator passes it as the pool `device`. Omitted / | |
| # bare "usb" ⇒ the first available Coral (single-Coral default). | |
| dev = (device or "").strip() | |
| opts = {"device": dev} if dev and dev not in ("usb", "all", "default") else None | |
| # The Coral is a SINGLE-CONTEXT USB device: only ONE process may hold the | |
| # libedgetpu delegate at a time. When a prior pool (e.g. a previous runner | |
| # generation, or a just-disposed device pool) is still releasing the USB, | |
| # a fresh load_delegate fails transiently. Retry with backoff instead of | |
| # hard-failing the whole pool — the device frees within a second or two of | |
| # the old holder exiting. (A genuinely-held device still fails after the | |
| # retries, with the real libedgetpu error surfaced.) | |
| import time as _time | |
| last_err = None | |
| delegate = None | |
| for attempt in range(6): | |
| try: | |
| delegate = load_delegate(lib_path, opts) if opts else load_delegate(lib_path) | |
| break | |
| except Exception as e: # noqa: BLE001 — surface the real error after retries | |
| last_err = e | |
| sys.stderr.write( | |
| f"EdgeTPU: delegate load attempt {attempt + 1}/6 failed " | |
| f"(device={dev or 'usb'}): {e}\n" | |
| ) | |
| sys.stderr.flush() | |
| _time.sleep(0.75 * (attempt + 1)) | |
| if delegate is None: | |
| raise ValueError( | |
| f"EdgeTPU: delegate not available after 6 attempts on device " | |
| f"'{dev or 'usb'}' — the Coral is held by another process or " | |
| f"absent. Last error: {last_err}" | |
| ) | |
| sys.stderr.write(f"EdgeTPU: delegate bound to device '{dev or 'usb'}'\n") | |
| _runtime_lib = { | |
| "Interpreter": Interpreter, | |
| "delegate": delegate, | |
| # invoke() on one interpreter is not re-entrant, and the EdgeTPU is a | |
| # single-context device — serialise all predicts through one lock. | |
| "lock": threading.Lock(), | |
| } | |
| sys.stderr.write(f"EdgeTPU: delegate loaded from {lib_path}\n") | |
| sys.stderr.flush() | |
| else: | |
| raise ValueError(f"Unknown runtime: {runtime}") | |
| # --------------------------------------------------------------------------- | |
| # Batch capability is a HARDWARE FACT, not a setting (D202) | |
| # --------------------------------------------------------------------------- | |
| # Whether a compiled graph may keep a DYNAMIC batch axis is decided here, from | |
| # the (engine, device) pair, at the moment of load — never per model, never by | |
| # an operator, never in a conversion script. The conversion scripts and the | |
| # model-studio convert path both emit dynamic-batch graphs (`ov.convert_model` | |
| # preserves the ONNX dynamic axis), and the same artifact is loaded on CPU, iGPU, | |
| # NPU and ANE nodes — so the artifact cannot know, and only the loader does. | |
| # | |
| # The one target that refuses: the Intel NPU. Its Level-Zero compiler rejects a | |
| # dynamic dimension outright ("Got negative shape dim bound: '-1'", then | |
| # ZE_RESULT_ERROR_INVALID_ARGUMENT). Found live 2026-08-20 on the hub, and the | |
| # symptom is the trap — a pool load failure is a RETRY, not an error the caller | |
| # sees, so an NPU-pinned node repeated "Loading step variant" forever and looked | |
| # busy rather than broken. Under AUTO it is quieter and just as wrong: the NPU | |
| # candidate silently loses every dynamic-batch model to the GPU fallback. | |
| # | |
| # This is COMPILE shape only. The pool's runtime request grouping (`batch_mode` | |
| # / `max_batch_size`, and the CoreML stacked-tensor fast path gated on the | |
| # export's own RangeDim) is a separate, untouched optimisation. | |
| STATIC_BATCH_PIN = 1 | |
| def _static_batch_required(runtime: str, device: str) -> bool: | |
| """THE authority: does this (engine, device) pair refuse a dynamic shape? | |
| Pure — no OpenVINO import, so it is unit-testable in a bare sandbox. | |
| openvino + NPU -> True (Level-Zero rejects any dynamic dim) | |
| openvino + CPU/GPU -> False | |
| onnxruntime -> False (every EP we ship handles dynamic axes) | |
| coreml -> False (a RangeDim export is exactly what the batched | |
| ANE dispatch path reads; pinning would delete it) | |
| edgetpu -> False (a tflite graph is static by construction) | |
| """ | |
| if runtime != "openvino": | |
| return False | |
| return device.strip().upper().startswith("NPU") | |
| def _static_batch_plan( | |
| shapes: "list[list[int | None] | None]", | |
| ) -> "tuple[dict[int, list[int]], list[int]]": | |
| """Pure. Given each input's dims (``None`` = dynamic axis, ``None`` for the | |
| whole entry = unreadable), return ``(plan, unpinnable)``: | |
| plan {input_index: static_dims} for inputs whose BATCH axis is | |
| dynamic and whose every other axis is already static. | |
| unpinnable [input_index] for inputs carrying a dynamic axis that is NOT | |
| the batch — there is no value we are allowed to invent for a | |
| spatial dim, so the caller compiles as-is and SAYS SO rather | |
| than letting the failure be silent. | |
| """ | |
| plan: "dict[int, list[int]]" = {} | |
| unpinnable: "list[int]" = [] | |
| for index, dims in enumerate(shapes): | |
| if dims is None: | |
| unpinnable.append(index) | |
| continue | |
| if len(dims) == 0: | |
| continue | |
| if any(d is None for d in dims[1:]): | |
| unpinnable.append(index) | |
| continue | |
| if dims[0] is not None: | |
| continue # already a static batch — the rule pins dynamic axes only | |
| plan[index] = [STATIC_BATCH_PIN] + [int(d) for d in dims[1:] if d is not None] | |
| return plan, unpinnable | |
| def _pin_batch_for_device(core: Any, source: Any, path: str, device: str) -> Any: | |
| """OpenVINO applicator of {@link _static_batch_required}. | |
| Returns what `compile_model` should be handed for `device`: the untouched | |
| `source` when the target accepts a dynamic shape (or the graph is already | |
| static), otherwise a model RESHAPED to a static batch of | |
| {@link STATIC_BATCH_PIN}. `source` is either the IR path (the plain load | |
| path) or an already-built model object (the PrePostProcessor path); the pin | |
| is applied to whichever, and it only ever touches axis 0, so a PPP model's | |
| rewritten uint8 NHWC input survives it. Every outcome is logged — the | |
| derivation must be readable in the pool stderr. | |
| (`_ov_dim_to_int` is defined further down with the other OV shape helpers.) | |
| """ | |
| name = os.path.basename(path) | |
| if not _static_batch_required("openvino", device): | |
| return source | |
| try: | |
| model = core.read_model(path) if isinstance(source, str) else source | |
| except Exception as exc: | |
| sys.stderr.write( | |
| f"OpenVINO: {device} needs a static shape but re-reading {name} failed " | |
| f"({exc}); compiling the source as-is\n" | |
| ) | |
| sys.stderr.flush() | |
| return source | |
| shapes: "list[list[int | None] | None]" = [] | |
| for inp in model.inputs: | |
| try: | |
| pshape = inp.get_partial_shape() | |
| shapes.append([_ov_dim_to_int(pshape[i]) for i in range(len(pshape))]) | |
| except Exception: | |
| shapes.append(None) | |
| plan, unpinnable = _static_batch_plan(shapes) | |
| if unpinnable: | |
| sys.stderr.write( | |
| f"OpenVINO: {name} has a DYNAMIC NON-BATCH axis on input(s) {unpinnable} " | |
| f"and {device} refuses dynamic dims — no value may be invented for a " | |
| f"spatial axis, so this compile is expected to fail\n" | |
| ) | |
| sys.stderr.flush() | |
| if not plan: | |
| sys.stderr.write( | |
| f"OpenVINO: batch shape for {name} on {device} — already static, no " | |
| f"reshape (derived from engine+device)\n" | |
| ) | |
| sys.stderr.flush() | |
| return source | |
| try: | |
| from openvino import PartialShape | |
| # `Model.reshape` mutates IN PLACE, and on the PrePostProcessor path | |
| # `source` IS the model object the compile loop reuses for the NEXT | |
| # candidate. Pinning it directly would leak this NPU-only decision into | |
| # the GPU/CPU fallback — a candidate that accepts a dynamic shape would | |
| # silently get the pinned graph. So reshape a CLONE and leave `source` | |
| # exactly as the caller built it. | |
| model = model.clone() if not isinstance(source, str) else model | |
| model.reshape({index: PartialShape(dims) for index, dims in plan.items()}) | |
| except Exception as exc: | |
| sys.stderr.write( | |
| f"OpenVINO: static-batch reshape of {name} for {device} failed ({exc}); " | |
| f"compiling the source as-is\n" | |
| ) | |
| sys.stderr.flush() | |
| return source | |
| sys.stderr.write( | |
| f"OpenVINO: batch shape for {name} on {device} — PINNED to " | |
| f"{STATIC_BATCH_PIN} on input(s) {sorted(plan)} because the target refuses " | |
| f"a dynamic dim (derived from engine+device, not configured)\n" | |
| ) | |
| sys.stderr.flush() | |
| return model | |
| def _load_model(slot: ModelSlot, config: dict) -> None: | |
| """Load a model into a slot using the active runtime. Thin adapter per backend.""" | |
| slot.config = dict(config) | |
| path = config["path"] | |
| if _runtime == "coreml": | |
| ct = _runtime_lib | |
| device = config.get("device", "all") | |
| compute_map = { | |
| "cpu": ct.ComputeUnit.CPU_ONLY, | |
| "gpu": ct.ComputeUnit.CPU_AND_GPU, | |
| "ane": ct.ComputeUnit.CPU_AND_NE, | |
| "all": ct.ComputeUnit.ALL, | |
| } | |
| # NOTE: a macOS behavior can cause the .mlpackage to be RECOMPILED on | |
| # every load with the compiled cache not reused until reboot. camstack | |
| # loads the .mlpackage on | |
| # every pool spawn, so it was exposed to the same trap. | |
| # _acquire_coreml_model routes through a persistent `.coreml-cache` | |
| # of the compiled `.mlmodelc` (CompiledMLModel on hit, plain MLModel | |
| # compile + persist on miss) and returns the spec read via | |
| # ct.utils.load_spec on hits — falling back to the plain | |
| # MLModel/get_spec() path on ANY cache failure. | |
| # | |
| # First, the package is loaded once in a CHILD process: a package that | |
| # aborts Core ML's native loader (SIGABRT, which Python cannot catch) | |
| # is refused here with `CoremlLoadCrashed` instead of killing this | |
| # worker and, through the respawn, every worker after it. | |
| import coreml_load_probe | |
| coreml_load_probe.ensure_coreml_loadable(path, device) | |
| model, spec = _acquire_coreml_model( | |
| ct, path, compute_map.get(device, ct.ComputeUnit.ALL), | |
| ) | |
| # Report the compute unit that was ACTUALLY applied so a silent | |
| # CPU-only placement is visible in the pool stderr (mirrors the | |
| # OpenVINO EXECUTION_DEVICES / ONNX provider reports). Both MLModel | |
| # and CompiledMLModel expose `compute_unit`. | |
| _actual_cu = getattr(model, "compute_unit", None) | |
| sys.stderr.write( | |
| f"CoreML: loaded {os.path.basename(path)} device={device} " | |
| f"compute_unit={_actual_cu}\n" | |
| ) | |
| sys.stderr.flush() | |
| input_spec = spec.description.input[0] | |
| slot.config["_input_name"] = input_spec.name | |
| slot.config["_input_type"] = input_spec.type.WhichOneof("Type") | |
| slot.config["_input_names"] = [i.name for i in spec.description.input] | |
| if slot.config["_input_type"] == "multiArrayType": | |
| slot.config["_input_shape"] = list(input_spec.type.multiArrayType.shape) | |
| # Detect batch flexibility — when the model's first axis is | |
| # declared as a `shapeRange` (RangeDim at export time), we | |
| # can stack N preprocessed frames into a single (N,3,H,W) | |
| # tensor and issue ONE predict call instead of iterating | |
| # the list-of-dicts path. One ANE dispatch per batch. | |
| mat = input_spec.type.multiArrayType | |
| if mat.WhichOneof("ShapeFlexibility") == "shapeRange" and len(mat.shapeRange.sizeRanges) > 0: | |
| batch_axis = mat.shapeRange.sizeRanges[0] | |
| if batch_axis.upperBound > 1: | |
| slot.config["_supports_batch"] = True | |
| slot.config["_max_batch"] = int(batch_axis.upperBound) | |
| slot.model = model | |
| slot.input_names = frozenset(i.name for i in spec.description.input) | |
| # Filter via slot.input_names atomically read from the CURRENT | |
| # slot. During a model `replace`, in-flight preprocess calls | |
| # may have added stale keys (iouThreshold, confidenceThreshold) | |
| # for the prior model that the new one rejects. | |
| def _predict(inp: dict, _slot: ModelSlot = slot) -> dict: | |
| names = _slot.input_names | |
| filtered = {k: v for k, v in inp.items() if k in names} if names else inp | |
| return _slot.model.predict(filtered) | |
| slot.predict_fn = _predict | |
| elif _runtime == "openvino": | |
| core = _runtime_lib | |
| ov_device_req = str(config.get("device", "AUTO")).upper() | |
| # Device switch (see _resolve_ov_devices above): an AUTO/default request | |
| # classifies the enumerated devices by FULL_DEVICE_NAME and yields an | |
| # ORDERED EXPLICIT candidate list (NPU>GPU>CPU) for compile-time | |
| # failover. An operator-pinned device (config "device" = gpu/cpu/npu) is | |
| # honored unchanged — compiled directly, no failover. | |
| if ov_device_req in ("AUTO", "DEFAULT", ""): | |
| available = list(core.available_devices) | |
| full_names: dict[str, str] = {} | |
| for dev in available: | |
| try: | |
| full_names[dev] = str(core.get_property(dev, "FULL_DEVICE_NAME")) | |
| except Exception: | |
| pass | |
| candidates = _resolve_ov_devices(available, full_names) | |
| sys.stderr.write( | |
| f"OpenVINO: explicit device candidates {candidates} " | |
| f"(available: {available}, names: {full_names})\n" | |
| ) | |
| sys.stderr.flush() | |
| else: | |
| candidates = [ov_device_req] # operator pin — compiled directly | |
| sys.stderr.write(f"OpenVINO: operator-pinned device {ov_device_req}\n") | |
| sys.stderr.flush() | |
| # PERFORMANCE_HINT is LATENCY: one | |
| # deterministic execution stream instead of N. THROUGHPUT spins up | |
| # multiple internal streams whose lazy request creation historically | |
| # collided with the AUTO handover AND multiplied memory (a factor in the | |
| # N100 OOM-freeze); it also RAISES OPTIMAL_NUMBER_OF_INFER_REQUESTS, which | |
| # enlarges the request pool at a memory / tail-latency cost. | |
| # LATENCY keeps a single stream; camera-level concurrency comes from the | |
| # fixed InferRequest pool built below. | |
| ov_config = {"PERFORMANCE_HINT": _OV_PERF_HINT} | |
| # Persist compiled GPU/NPU kernels so the multi-second JIT compile is paid | |
| # once (first boot) instead of on every pool spawn. The cold compile storm | |
| # (~25s for the full model set on an N100 iGPU) is where frames pile up; | |
| # caching the kernels shrinks that window on subsequent boots. | |
| try: | |
| _cache_dir = os.path.join(os.path.dirname(os.path.abspath(path)), ".ov-cache") | |
| os.makedirs(_cache_dir, exist_ok=True) | |
| ov_config["CACHE_DIR"] = _cache_dir | |
| except OSError: | |
| pass | |
| # Fix 2 — fold uint8→float/scale/layout preprocessing INTO the compiled | |
| # graph via PrePostProcessor. When it succeeds we compile the PPP model | |
| # object (fed uint8 NHWC at inference time); on ANY failure we fall back | |
| # to compiling the plain `path` (byte-identical to the historical float | |
| # preprocess). Always on for OpenVINO (tests flip `_OV_PPP_ENABLED`). | |
| ppp_meta: Optional[dict] = None | |
| compile_source: Any = path | |
| if _OV_PPP_ENABLED: | |
| try: | |
| ppp_model, ppp_meta = _build_ov_ppp_model(core, path, config) | |
| if ppp_model is not None: | |
| compile_source = ppp_model | |
| sys.stderr.write( | |
| f"OpenVINO: PrePostProcessor folded for {os.path.basename(path)} " | |
| f"(uint8 NHWC in, {ppp_meta}); Python preprocess stays uint8\n" | |
| ) | |
| sys.stderr.flush() | |
| except Exception as exc: | |
| ppp_meta = None | |
| compile_source = path | |
| sys.stderr.write( | |
| f"OpenVINO: PrePostProcessor setup failed for {os.path.basename(path)} " | |
| f"({exc}); falling back to Python float preprocess\n" | |
| ) | |
| sys.stderr.flush() | |
| # Compile-time failover over the explicit candidate list (NO AUTO): try | |
| # each device in order (NPU>GPU>CPU) and stop at the first that compiles. | |
| # CPU is the guaranteed floor. If every candidate fails, re-raise — the | |
| # existing load/replace error handler reports the failure. | |
| compiled = None | |
| ov_device = None | |
| # What the compilation bound actually became on the device that won — | |
| # 0 means the device does not expose the key, so this compile spent | |
| # `nproc` threads for the life of the process and the operator must be | |
| # able to SEE that rather than assume the bound applied. | |
| ov_compile_threads = 0 | |
| ov_precision = "device-default" | |
| last_exc: Optional[BaseException] = None | |
| for candidate in candidates: | |
| try: | |
| # Batch capability is derived from THIS candidate, not from the | |
| # artifact and not from a setting (D202): a graph with a dynamic | |
| # batch axis is reshaped static before it can reach a compiler | |
| # that refuses dynamic dims (the Intel NPU), and is handed over | |
| # untouched to every other target. | |
| candidate_source = _pin_batch_for_device( | |
| core, compile_source, path, candidate, | |
| ) | |
| # Bound the per-compile TBB pool this candidate would otherwise | |
| # size to `nproc` and never give back (see | |
| # OV_COMPILATION_NUM_THREADS). Gated on the device's own | |
| # SUPPORTED_PROPERTIES so a plugin that does not know the key | |
| # compiles unchanged instead of failing over to the CPU. | |
| candidate_supported = _ov_supported_properties(core, candidate) | |
| candidate_config = { | |
| **ov_config, | |
| **_ov_compilation_properties( | |
| candidate_supported, | |
| OV_COMPILATION_NUM_THREADS, | |
| ), | |
| # A model the GPU must not run at its default f16 | |
| # (catalog `gpuInferencePrecision`). | |
| **_ov_precision_properties(config, candidate, candidate_supported), | |
| } | |
| compiled = core.compile_model( | |
| candidate_source, device_name=candidate, config=candidate_config, | |
| ) | |
| ov_device = candidate | |
| ov_compile_threads = int( | |
| candidate_config.get("COMPILATION_NUM_THREADS", 0) | |
| ) | |
| ov_precision = str( | |
| candidate_config.get("INFERENCE_PRECISION_HINT", "device-default") | |
| ) | |
| break | |
| except Exception as exc: | |
| last_exc = exc | |
| sys.stderr.write(f"OpenVINO: compile on {candidate} failed: {exc}\n") | |
| if candidate != candidates[-1]: | |
| sys.stderr.write("OpenVINO: falling back to next device.\n") | |
| sys.stderr.flush() | |
| if compiled is None: | |
| if last_exc is not None: | |
| raise last_exc | |
| raise RuntimeError("OpenVINO: no device available to compile the model") | |
| # Name the resolved mode + what OpenVINO actually placed the model on | |
| # (the EXECUTION_DEVICES property) so a silent CPU | |
| # placement is visible in the pool stderr. | |
| try: | |
| _exec_devices = compiled.get_property("EXECUTION_DEVICES") | |
| except Exception: | |
| _exec_devices = "unknown" | |
| sys.stderr.write( | |
| f"OpenVINO: compiled {os.path.basename(path)} mode={ov_device} " | |
| f"EXECUTION_DEVICES={_exec_devices} compilationThreads=" | |
| f"{ov_compile_threads if ov_compile_threads > 0 else 'unbounded(device rejects the key)'}" | |
| f" inferencePrecision={ov_precision}\n" | |
| ) | |
| # A model that needs a GPU precision the device would not take runs at | |
| # f16 and may return garbage — say so rather than look healthy. | |
| if ( | |
| config.get("gpuInferencePrecision") in _OV_GPU_PRECISIONS | |
| and str(ov_device).upper().startswith("GPU") | |
| and ov_precision == "device-default" | |
| ): | |
| sys.stderr.write( | |
| f"OpenVINO: WARNING {os.path.basename(path)} requires " | |
| f"inferencePrecision={config.get('gpuInferencePrecision')} on {ov_device}, " | |
| "which does not list INFERENCE_PRECISION_HINT — it runs at the device default " | |
| "and its detections may be invalid\n" | |
| ) | |
| sys.stderr.flush() | |
| # An OpenVINO IR converted from ONNX can carry an UNNAMED output tensor | |
| # (e.g. `yolov8n-package`): `get_any_name()` then raises | |
| # "Attempt to get a name for a Tensor without names", which previously | |
| # failed the whole model load on OpenVINO nodes (the package-detection | |
| # step could never run on OpenVINO). Outputs are fetched POSITIONALLY in | |
| # `predict` below, so a synthetic key per unnamed port is sufficient. | |
| output_names = [] | |
| for _i, _o in enumerate(compiled.outputs): | |
| try: | |
| output_names.append(_o.get_any_name()) | |
| except Exception: | |
| output_names.append(f"output{_i}") | |
| # Record the device's optimal infer-request count so the dispatcher | |
| # can size the predict pool to actually feed the streams. | |
| global _ov_optimal_reqs | |
| try: | |
| opt = int(compiled.get_property("OPTIMAL_NUMBER_OF_INFER_REQUESTS")) | |
| if opt > _ov_optimal_reqs: | |
| _ov_optimal_reqs = opt | |
| except Exception: | |
| pass | |
| # Size N ONCE from the device's optimal infer-request count (falling back | |
| # to OV_DEFAULT_CONCURRENCY only when the property is missing/zero) so it | |
| # matches the predict-pool worker count sized in _run(); `_ov_request_count` | |
| # is the one place it is computed. | |
| import queue as _queue | |
| try: | |
| _n_req = int(compiled.get_property("OPTIMAL_NUMBER_OF_INFER_REQUESTS")) | |
| except Exception: | |
| _n_req = 0 | |
| _n_req = _ov_request_count(_n_req) | |
| # Create ALL InferRequests ONCE, now, post-compile. With AUTO gone there is no CPU->GPU handover, so no request | |
| # can be born mid-handover and bind to stale ports — the score_8 race | |
| # is structurally removed. Requests are handed out through a thread- | |
| # safe queue: an InferRequest is never touched by two threads at once | |
| # (checkout/return), so sharing a fixed pool is safe and avoids the old | |
| # thread-local lazy creation. | |
| _req_pool: "_queue.Queue" = _queue.Queue() | |
| for _ in range(_n_req): | |
| _req_pool.put(compiled.create_infer_request()) | |
| sys.stderr.write(f"OpenVINO: created {_n_req} InferRequest(s) up-front on {ov_device}\n") | |
| sys.stderr.flush() | |
| def predict(inp_dict: dict, _names=output_names, _pool=_req_pool) -> dict: | |
| inp = list(inp_dict.values())[0] | |
| req = _pool.get() | |
| try: | |
| result = req.infer(inp) | |
| finally: | |
| _pool.put(req) | |
| # Output fetch by positional index — stable and harmless now that | |
| # the ports no longer diverge (single explicit device, no handover). | |
| return {name: result[i] for i, name in enumerate(_names)} | |
| slot.model = compiled | |
| slot.predict_fn = predict | |
| # When the PPP preprocessor is folded in, `_preprocess` produces uint8 | |
| # NHWC and short-circuits BEFORE the NHWC/NCHW float logic — so the | |
| # `_input_shape` heuristic below is irrelevant for PPP models. Record | |
| # the plan so `_preprocess` takes the uint8 fast path. | |
| if ppp_meta is not None: | |
| slot.config["_ov_ppp"] = ppp_meta | |
| else: | |
| # Capture the model's declared input layout so `_preprocess` feeds | |
| # NHWC vs NCHW correctly (ArcFace/CLIP are NHWC). Best-effort — a | |
| # failure leaves `_input_shape` unset → NCHW fallback (unchanged). | |
| try: | |
| slot.config["_input_shape"] = list(compiled.inputs[0].get_partial_shape()) | |
| except Exception: | |
| pass | |
| elif _runtime == "onnxruntime": | |
| ort = _runtime_lib | |
| ort_device = str(config.get("device", "") or "") | |
| # Ordered providers — hardware EP first, CPU LAST as | |
| # fallback (see _resolve_onnx_providers): a "cuda"/"coreml" pin puts | |
| # that EP first, a "cpu" pin compiles CPU only, and the AUTO/default | |
| # case derives the hardware EP from the platform. | |
| try: | |
| available_eps = list(ort.get_available_providers()) | |
| except Exception: | |
| available_eps = [ONNX_CPU_EP] | |
| providers = _resolve_onnx_providers( | |
| ort_device, sys.platform, platform.machine(), available_eps, | |
| ) | |
| session = ort.InferenceSession(path, providers=providers) | |
| # Report the EP that ACTUALLY initialized (CPU stripped from the log) | |
| # so a silent CPU fallback is visible. | |
| _active_eps = [p for p in session.get_providers() if p != ONNX_CPU_EP] | |
| _requested = [p[0] if isinstance(p, tuple) else p for p in providers] | |
| sys.stderr.write( | |
| f"ONNX: execution provider {', '.join(_active_eps) or ONNX_CPU_EP} " | |
| f"(device={ort_device or 'auto'}, requested={_requested}, " | |
| f"available={available_eps})\n" | |
| ) | |
| sys.stderr.flush() | |
| output_names = [o.name for o in session.get_outputs()] | |
| slot.config["_input_name"] = session.get_inputs()[0].name | |
| # Declared input layout (ArcFace/CLIP are NHWC) → correct preprocess. | |
| try: | |
| slot.config["_input_shape"] = session.get_inputs()[0].shape | |
| except Exception: | |
| pass | |
| # Declared input DTYPE. super-gradients / Frigate YOLO-NAS exports bake | |
| # their preprocessing into the graph and take `tensor(uint8)` — the | |
| # preprocess must feed raw uint8 pixels, not the float32 [0,1] tensor | |
| # (a hard onnxruntime type error otherwise). | |
| try: | |
| if session.get_inputs()[0].type == "tensor(uint8)": | |
| slot.config["_input_dtype"] = "uint8" | |
| except Exception: | |
| pass | |
| def predict(inp_dict: dict) -> dict: | |
| outputs = session.run(output_names, inp_dict) | |
| return {name: out for name, out in zip(output_names, outputs)} | |
| slot.model = session | |
| slot.predict_fn = predict | |
| elif _runtime == "edgetpu": | |
| Interpreter = _runtime_lib["Interpreter"] | |
| delegate = _runtime_lib["delegate"] | |
| lock = _runtime_lib["lock"] | |
| interp = Interpreter(path, experimental_delegates=[delegate]) | |
| interp.allocate_tensors() | |
| in_det = interp.get_input_details()[0] | |
| out_dets = interp.get_output_details() | |
| _, in_h, in_w, _ = (int(x) for x in in_det["shape"]) | |
| in_dtype = np.dtype(in_det["dtype"]).name # 'uint8' | 'int8' | 'float32' | |
| slot.config["_input_name"] = in_det["name"] | |
| slot.config["_edgetpu_in_dtype"] = in_dtype | |
| slot.config["inputWidth"] = in_w | |
| slot.config["inputHeight"] = in_h | |
| # The FILE's input size is the truth for a compiled Edge TPU graph: a | |
| # Coral build may be compiled at a smaller input than the catalog | |
| # entry's other formats (activations share the 8 MB on-chip memory | |
| # with the parameter cache, D710), and a decoder sized from the | |
| # catalog would put every prior in the wrong cell. | |
| slot.config["inputSize"] = max(in_w, in_h) | |
| in_index = int(in_det["index"]) | |
| out_index = [int(o["index"]) for o in out_dets] | |
| from edgetpu_dequant import edgetpu_output_keys | |
| out_keys = edgetpu_output_keys( | |
| [str(o.get("name") or "") for o in out_dets], slot.config.get("postprocessor") | |
| ) | |
| sys.stderr.write( | |
| f"EdgeTPU: loaded {os.path.basename(path)} in={in_w}x{in_h} " | |
| f"dtype={in_dtype} outputs={len(out_dets)}\n" | |
| ) | |
| sys.stderr.flush() | |
| def predict( | |
| inp_dict: dict, | |
| _interp=interp, _lock=lock, _in=in_index, _outs=out_index, _details=out_dets, | |
| _keys=out_keys, | |
| ) -> dict: | |
| from edgetpu_dequant import dequant_edgetpu_output | |
| tensor = next(iter(inp_dict.values())) | |
| with _lock: | |
| _interp.set_tensor(_in, tensor) | |
| _interp.invoke() | |
| # `get_tensor` returns a copy, so the Coral is released before | |
| # the dequantisation below runs. | |
| raws = [_interp.get_tensor(idx) for idx in _outs] | |
| # Positional keys — ssd reads "0".."3"; yolo-split-dfl pairs by | |
| # shape — except for a postprocessor that matches by name | |
| # (`edgetpu_output_keys`). Integer outputs are dequantised here so | |
| # every postprocessor sees float32 (SSD Detection_PostProcess is | |
| # already float and passes through). | |
| result = {} | |
| for i, raw in enumerate(raws): | |
| detail = _details[i] | |
| q = detail.get("quantization") or (0.0, 0) | |
| scale = float(q[0]) if q else 0.0 | |
| zp = int(q[1]) if q else 0 | |
| dtype_name = str(np.dtype(detail["dtype"]).name) | |
| result[_keys[i]] = dequant_edgetpu_output( | |
| raw, dtype_name=dtype_name, scale=scale, zero_point=zp | |
| ) | |
| return result | |
| slot.model = interp | |
| slot.predict_fn = predict | |
| slot.loaded = True | |
| def _unload_model(slot: ModelSlot) -> None: | |
| slot.model = None | |
| slot.predict_fn = None | |
| slot.input_names = frozenset() | |
| slot.config = {} | |
| slot.loaded = False | |
| # --------------------------------------------------------------------------- | |
| # Preprocess + run_inference — pure sync helpers called from worker threads | |
| # --------------------------------------------------------------------------- | |
| # Preprocess cache for bench frames only. Keyed by (_bench_frame_id, inputSize) | |
| # where _bench_frame_id is a stable integer tag set on PIL Images stored in | |
| # frame_cache. Camera frames don't have this tag → no cache → correct behavior. | |
| # Avoids the CPython id() reuse problem that caused stale detections. | |
| _bench_preprocess_cache: dict[tuple[int, int], tuple[dict, float, tuple[int, int]]] = {} | |
| # Verifiable proof that the bench preprocess cache is actually HIT on the | |
| # sustained-throughput BATCH path. A tagged frame that misses runs the FULL | |
| # decode-independent preprocess (letterbox + float/255 + transpose); a hit | |
| # returns the pre-built input tensor untouched — pure inference, no CPU | |
| # preprocess. The historical bug was that MSG_INFER_BATCH wrapped each item | |
| # fresh WITHOUT the tag, so the cache was NEVER hit on the batch path and every | |
| # inference re-ran the full preprocess. The sustained benchmark now asserts | |
| # hits ≫ misses (surfaced in the `status` command) so that regression is caught. | |
| _bench_cache_hits: int = 0 | |
| _bench_cache_misses: int = 0 | |
| def _bench_cache_stats() -> dict: | |
| """(hits, misses) for the bench preprocess cache — surfaced in `status`.""" | |
| return {"hits": _bench_cache_hits, "misses": _bench_cache_misses} | |
| def _is_channels_last(input_shape) -> bool: | |
| """True when a 4-D model input is NHWC ([N,H,W,C]) rather than NCHW | |
| ([N,C,H,W]). Heuristic: the LAST dim is a channel count (1 or 3) and the | |
| SECOND dim is a concrete non-channel value. ArcFace / CLIP declare | |
| [N,112,112,3] → True; standard detectors [N,3,H,W] → False. Robust to | |
| dynamic/None dims → returns False (NCHW, the prior unconditional | |
| behaviour) whenever it can't be sure, so detectors are never touched.""" | |
| if input_shape is None or len(input_shape) != 4: | |
| return False | |
| def _as_int(v) -> "int | None": | |
| try: | |
| return int(v) | |
| except (TypeError, ValueError): | |
| pass | |
| # OpenVINO `Dimension` objects don't support int() (TypeError even | |
| # for STATIC dims — verified live on the hub: arcface's captured | |
| # shape [?,112,112,3] read as all-None → NCHW fallback → NPU shape | |
| # error). Static Dimensions expose get_length(); dynamic stay None. | |
| try: | |
| if getattr(v, "is_static", False): | |
| return int(v.get_length()) | |
| except (TypeError, ValueError, AttributeError): | |
| pass | |
| return None | |
| last = _as_int(input_shape[-1]) | |
| second = _as_int(input_shape[1]) | |
| return last in (1, 3) and second is not None and second not in (1, 3) | |
| def _resolve_input_dims(config: dict) -> tuple[int, int, int]: | |
| """Resolve (width, height, channels) for the model input. | |
| Explicit `inputWidth`/`inputHeight` (threaded from the catalog's | |
| `inputSize.{width,height}`) win; they fall back to the square `inputSize` | |
| (the max side) for legacy square models. `inputChannels` defaults to 3 | |
| (RGB); 1 selects the grayscale CTC recognizer path. Pure — no numpy/PIL. | |
| """ | |
| input_size = int(config.get("inputSize", 640)) | |
| width = int(config.get("inputWidth") or input_size) | |
| height = int(config.get("inputHeight") or input_size) | |
| channels = int(config.get("inputChannels") or 3) | |
| return width, height, channels | |
| def _to_model_tensor(arr: "np.ndarray", channels_last: bool) -> "np.ndarray": | |
| """Build the batched RGB model-input tensor from an HWC float array in [0,1]. | |
| - ``channels_last`` (ArcFace/CLIP NHWC) → ``[1,H,W,C]`` untransposed. | |
| - otherwise → standard detector NCHW ``[1,C,H,W]``. | |
| Grayscale CTC recognizers do NOT use this — they need aspect-preserving | |
| resize + edge padding, see {@link _preprocess_ctc_gray}. Pure numpy. | |
| """ | |
| if channels_last: | |
| return arr[np.newaxis].astype(np.float32) | |
| return arr.transpose(2, 0, 1)[np.newaxis].astype(np.float32) | |
| _IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) | |
| _IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) | |
| # InsightFace's SCRFD input contract: upstream feeds `(pixel - 127.5) / 128` | |
| # ([-0.996, 0.996]), not the pool's historical plain `/255` ([0,1]). Expressed | |
| # against the pool's already-`/255` HWC array (`arr`, not raw pixels): | |
| # (arr*255 - 127.5) / 128 == (arr - _SCRFD_MEAN) / _SCRFD_STD | |
| # with mean/std pre-divided by 255, the same shape of transform as the | |
| # ImageNet branch below. See `postprocessors/scrfd.py` for the upstream | |
| # citation and the measured effect (188 → 196 GT faces found, | |
| # 2026-09-26 model-replacement spike, §2.4-2). | |
| _SCRFD_MEAN = 127.5 / 255.0 | |
| _SCRFD_STD = 128.0 / 255.0 | |
| def _apply_normalization(arr: "np.ndarray", config: dict) -> "np.ndarray": | |
| """Apply per-model input normalization to an HWC RGB float array in [0,1]. | |
| ``inputNormalization == 'imagenet'`` subtracts the ImageNet per-channel | |
| mean and divides by the std (RGB), matching the EfficientNet-Lite0 / | |
| MobileNetV3 animal + vehicle classifiers whose ``labels.json`` declares | |
| ``normalize: imagenet``. ``'scrfd'`` applies InsightFace's SCRFD contract | |
| (`_SCRFD_MEAN`/`_SCRFD_STD` above) — the `scrfd-2.5g` catalog entry. | |
| Absent / ``'none'`` / ``'zero-one'`` keep the plain ``/255`` array | |
| byte-identical to the historical unconditional path (most detectors, | |
| CLIP/ArcFace, and the AIY bird classifier which bakes its own scale). Runs | |
| on the HWC array BEFORE the optional NHWC/NCHW transpose in | |
| :func:`_to_model_tensor`, so it is layout-agnostic (channel axis is always | |
| last here). Pure numpy. | |
| """ | |
| mode = config.get("inputNormalization") | |
| if mode not in ("imagenet", "scrfd"): | |
| return arr | |
| if arr.ndim != 3 or arr.shape[-1] != 3: | |
| return arr | |
| if mode == "scrfd": | |
| return (arr - _SCRFD_MEAN) / _SCRFD_STD | |
| return (arr - _IMAGENET_MEAN) / _IMAGENET_STD | |
| def _preprocess_ctc_gray(img: "Image.Image", input_w: int, input_h: int) -> "np.ndarray": | |
| """EasyOCR-accurate preprocess for a fixed-width grayscale CTC recognizer. | |
| Mirrors EasyOCR's AlignCollate + NormalizePAD exactly: | |
| 1. grayscale by luminance (PIL ``convert('L')``), | |
| 2. resize to height ``input_h`` PRESERVING aspect ratio (width capped at | |
| ``input_w``) with BICUBIC — the recognizer was trained on undistorted | |
| glyphs; stretching the crop to the full ``input_w`` warps glyph aspect | |
| ~4x and the model reads noise (low-confidence garbage), | |
| 3. normalize to ``[-1,1]`` (``(x/255 - 0.5)/0.5``), | |
| 4. right-pad the width up to ``input_w`` by replicating the last column | |
| (EasyOCR NormalizePAD edge padding — NOT zeros). | |
| Returns a ``[1,1,input_h,input_w]`` float32 tensor. Pure numpy/PIL. | |
| """ | |
| gray = img.convert("L") | |
| ow, oh = gray.size | |
| ratio = (ow / oh) if oh > 0 else 1.0 | |
| resized_w = int(np.ceil(input_h * ratio)) | |
| resized_w = max(1, min(input_w, resized_w)) | |
| gray = gray.resize((resized_w, input_h), Image.BICUBIC) | |
| arr = np.asarray(gray, dtype=np.float32) / 255.0 # (H, resized_w) | |
| arr = (arr - 0.5) / 0.5 # [-1, 1] | |
| if resized_w < input_w: | |
| edge = arr[:, resized_w - 1 : resized_w] | |
| pad = np.repeat(edge, input_w - resized_w, axis=1) | |
| arr = np.concatenate([arr, pad], axis=1) | |
| return arr[np.newaxis, np.newaxis].astype(np.float32) # (1, 1, H, W) | |
| # --------------------------------------------------------------------------- | |
| # Fix 2 — OpenVINO PrePostProcessor (preprocess folded into the compiled graph) | |
| # --------------------------------------------------------------------------- | |
| # Always ON for OpenVINO; a module constant only so tests can exercise the | |
| # plain-path compile. A failed fold already falls back per model (above). | |
| _OV_PPP_ENABLED: bool = True | |
| def _ov_dim_to_int(dim: Any) -> "int | None": | |
| """Best-effort static-length extraction from an OpenVINO Dimension (or a | |
| plain int). Dynamic / unknown dims return None. Pure — no numpy.""" | |
| try: | |
| return int(dim) | |
| except (TypeError, ValueError): | |
| pass | |
| try: | |
| if getattr(dim, "is_static", False): | |
| return int(dim.get_length()) | |
| except (TypeError, ValueError, AttributeError): | |
| pass | |
| return None | |
| def _ov_ppp_plan( | |
| dims: "list[int | None] | None", | |
| input_channels: int, | |
| preprocess_mode: str, | |
| normalization: "str | None", | |
| enabled: bool, | |
| ) -> "dict | None": | |
| """Pure decision: should preprocessing be folded into the OpenVINO graph? | |
| `dims` is the model's declared input partial-shape as a list of int|None in | |
| NCHW order. We only fold for standard 3-channel NCHW detectors/classifiers | |
| with a STATIC spatial size — grayscale CTC recognizers (channels==1) and | |
| NHWC-native embedders (ArcFace/CLIP, channel dim last) keep the pure-Python | |
| path. Returns a `ppp_meta` dict describing the Python-side uint8 producer, | |
| or None when PPP is not applicable. No OpenVINO / numpy imports — | |
| unit-testable in isolation.""" | |
| if not enabled: | |
| return None | |
| if input_channels == 1: | |
| return None | |
| if dims is None or len(dims) != 4: | |
| return None | |
| _n, c, h, w = dims | |
| # NCHW with a concrete 3-channel dim. A last-dim of 3 (NHWC-native) leaves | |
| # c as a spatial value → rejected here, so embedders stay on the float path. | |
| if c != 3: | |
| return None | |
| if h is None or w is None or h <= 0 or w <= 0: | |
| return None | |
| return { | |
| "inputW": int(w), | |
| "inputH": int(h), | |
| "size": int(max(w, h)), | |
| "letterbox": preprocess_mode == "letterbox", | |
| "normalization": normalization or "none", | |
| } | |
| def _preprocess_ov_uint8( | |
| img: "Image.Image", plan: dict, | |
| ) -> tuple["np.ndarray", float, tuple[int, int]]: | |
| """Produce the uint8 NHWC ``[1,H,W,3]`` input for a PPP-folded OV model. | |
| Letterbox (aspect-preserving pad) or plain resize stays in cheap uint8 PIL | |
| space; the compiled OV graph does the float convert + /255 + optional | |
| mean/std + NHWC→NCHW transpose on the inference device. Returns | |
| ``(tensor, scale, pad)`` — the scale/pad still drive the postprocessor's box | |
| mapping exactly like the letterbox float path. Pure numpy/PIL.""" | |
| if plan["letterbox"]: | |
| canvas, scale_val, pad = letterbox_image(img, plan["size"]) | |
| else: | |
| w, h = plan["inputW"], plan["inputH"] | |
| if img.size == (w, h): | |
| canvas, scale_val, pad = img, 1.0, (0, 0) | |
| else: | |
| canvas, scale_val, pad = img.resize((w, h), Image.BILINEAR), 1.0, (0, 0) | |
| pix = np.asarray(canvas, dtype=np.uint8) | |
| if pix.ndim == 2: | |
| pix = np.stack([pix] * 3, axis=-1) | |
| return pix[np.newaxis], scale_val, pad | |
| def _build_ov_ppp_model(core: Any, path: str, config: dict) -> "tuple[Any, dict] | tuple[None, None]": | |
| """Read the OpenVINO IR and, when eligible, attach a PrePostProcessor that | |
| accepts a uint8 NHWC tensor and folds convert→scale(→mean/std)→layout into | |
| the compiled graph. Returns ``(model, ppp_meta)`` to compile, or | |
| ``(None, None)`` when PPP is not applied (caller compiles the path directly, | |
| byte-identical to the historical float path). Raises on hard OV errors so | |
| the caller's try/except falls back safely.""" | |
| from openvino import Layout, Type | |
| from openvino.preprocess import PrePostProcessor | |
| model = core.read_model(path) | |
| inp = model.input() | |
| try: | |
| pshape = inp.get_partial_shape() | |
| dims = [_ov_dim_to_int(pshape[i]) for i in range(len(pshape))] | |
| except Exception: | |
| dims = None | |
| _iw, _ih, input_channels = _resolve_input_dims(config) | |
| plan = _ov_ppp_plan( | |
| dims, | |
| input_channels, | |
| config.get("preprocessMode", "letterbox"), | |
| config.get("inputNormalization"), | |
| _OV_PPP_ENABLED, | |
| ) | |
| if plan is None: | |
| return None, None | |
| ppp = PrePostProcessor(model) | |
| # Tensor as the caller feeds it: uint8, NHWC, already at the model's spatial | |
| # size (Python letterboxes/resizes to size in uint8) → OV needs no resize. | |
| ppp.input().tensor().set_element_type(Type.u8).set_layout(Layout("NHWC")) | |
| ppp.input().model().set_layout(Layout("NCHW")) | |
| steps = ppp.input().preprocess().convert_element_type(Type.f32) | |
| if plan["normalization"] == "imagenet": | |
| # ((x/255) - mean) / std, applied in add-order (scale divides, mean | |
| # subtracts). RGB channel order matches the fed PIL tensor. | |
| steps.scale([255.0, 255.0, 255.0]) | |
| steps.mean([float(v) for v in _IMAGENET_MEAN]) | |
| steps.scale([float(v) for v in _IMAGENET_STD]) | |
| elif plan["normalization"] == "scrfd": | |
| # ((x/255) - _SCRFD_MEAN) / _SCRFD_STD — InsightFace's SCRFD contract, | |
| # folded the same way as the ImageNet branch above (see | |
| # `_apply_normalization`/`_SCRFD_MEAN`/`_SCRFD_STD`). | |
| steps.scale([255.0, 255.0, 255.0]) | |
| steps.mean([_SCRFD_MEAN, _SCRFD_MEAN, _SCRFD_MEAN]) | |
| steps.scale([_SCRFD_STD, _SCRFD_STD, _SCRFD_STD]) | |
| else: | |
| steps.scale([255.0, 255.0, 255.0]) | |
| built = ppp.build() | |
| return built, plan | |
| def _preprocess(img: Image.Image, config: dict) -> tuple[dict, float, tuple[int, int]]: | |
| input_size = config.get("inputSize", 640) | |
| input_w, input_h, input_channels = _resolve_input_dims(config) | |
| # Bench frames have _bench_frame_id tag → use preprocess cache | |
| bench_fid = getattr(img, '_bench_frame_id', None) | |
| if bench_fid is not None: | |
| global _bench_cache_hits, _bench_cache_misses | |
| cache_key = (bench_fid, input_size) | |
| cached = _bench_preprocess_cache.get(cache_key) | |
| if cached is not None: | |
| _bench_cache_hits += 1 | |
| return cached | |
| _bench_cache_misses += 1 | |
| preprocess_mode = config.get("preprocessMode", "letterbox") | |
| input_dict: dict = {} | |
| if _runtime == "coreml": | |
| input_name = config.get("_input_name", "image") | |
| input_type = config.get("_input_type", "imageType") | |
| input_shape = config.get("_input_shape") | |
| if input_type == "imageType": | |
| # FAST PATH — stay in uint8 PIL space. CoreML's predict() | |
| # accepts a PIL Image directly and handles the | |
| # BGR/normalize/CVPixelBuffer dance internally on its own | |
| # (zero-copy when the size matches the model input). Going | |
| # via numpy float32 → /255 → ×255 → astype(uint8) → Image. | |
| # fromarray() costs ~16ms per 640×640 frame on M3 Pro and is | |
| # PURE WASTE for imageType inputs. | |
| if preprocess_mode == "letterbox": | |
| canvas, scale_val, pad = letterbox_image(img, input_size) | |
| elif img.size == (input_size, input_size): | |
| canvas, scale_val, pad = img, 1.0, (0, 0) | |
| else: | |
| canvas, scale_val, pad = img.resize((input_size, input_size), Image.BILINEAR), 1.0, (0, 0) | |
| input_dict[input_name] = canvas | |
| arr = None # not used downstream for imageType — postprocess runs on predictions | |
| elif input_shape is not None and len(input_shape) == 4: | |
| if input_channels == 1: | |
| # Grayscale CTC recognizer — EasyOCR aspect-preserving preprocess. | |
| input_arr = _preprocess_ctc_gray(img, input_w, input_h) | |
| scale_val = 1.0 | |
| pad = (0, 0) | |
| else: | |
| if preprocess_mode == "letterbox": | |
| arr, scale_val, pad = letterbox(img, input_size) | |
| else: | |
| arr = resize_image(img, input_w, input_h) | |
| scale_val = 1.0 | |
| pad = (0, 0) | |
| # Per-model input normalization (imagenet mean/std for the | |
| # animal + vehicle classifiers; no-op for everything else). | |
| arr = _apply_normalization(arr, config) | |
| # NHWC vs NCHW from the model's declared last dim (1/3 ⇒ | |
| # channels-last), exactly as before for 3-channel models. | |
| _, _, _, w_or_c = input_shape | |
| input_arr = _to_model_tensor(arr, w_or_c in (1, 3)) | |
| input_dict[input_name] = input_arr | |
| else: | |
| input_arr = arr.transpose(2, 0, 1)[np.newaxis].astype(np.float32) | |
| input_dict[input_name] = input_arr | |
| input_names = config.get("_input_names", []) | |
| if "iouThreshold" in input_names: | |
| input_dict["iouThreshold"] = 0.45 | |
| if "confidenceThreshold" in input_names: | |
| input_dict["confidenceThreshold"] = config.get("confidence", 0.25) | |
| elif _runtime == "edgetpu": | |
| # EdgeTPU models are FULL-INTEGER quantized: feed raw uint8 (or int8) | |
| # pixels in NHWC [1,H,W,C] — NO float /255 normalization (that is baked | |
| # into the model's quantization params). | |
| input_name = config.get("_input_name", "images") | |
| in_w = int(config.get("inputWidth") or config.get("inputSize", 300)) | |
| in_h = int(config.get("inputHeight") or config.get("inputSize", 300)) | |
| if preprocess_mode == "letterbox": | |
| canvas, scale_val, pad = letterbox_image(img, max(in_w, in_h)) | |
| rgb = canvas | |
| else: | |
| rgb = img if img.size == (in_w, in_h) else img.resize((in_w, in_h), Image.BILINEAR) | |
| scale_val, pad = 1.0, (0, 0) | |
| pix = np.asarray(rgb, dtype=np.uint8) | |
| if pix.ndim == 2: | |
| pix = np.stack([pix] * 3, axis=-1) | |
| if config.get("_edgetpu_in_dtype") == "int8": | |
| input_dict[input_name] = (pix.astype(np.int16) - 128).astype(np.int8)[np.newaxis] | |
| else: | |
| input_dict[input_name] = pix[np.newaxis] # uint8 NHWC | |
| else: | |
| # OpenVINO / ONNX — always need a float tensor. | |
| input_name = config.get("_input_name", "images") | |
| # Fix 2 — OpenVINO PrePostProcessor fast path. When the model was | |
| # compiled with a folded uint8→float/scale/layout preprocessor | |
| # (`_ov_ppp` plan present), skip ALL numpy float work here: produce a | |
| # uint8 NHWC tensor (aspect-preserving letterbox stays in cheap PIL | |
| # space) and let the compiled OV graph do convert/scale/transpose on | |
| # the inference device. Removes the per-frame /255 + transpose that | |
| # dominated the non-inference CPU cost in production. | |
| ov_ppp = config.get("_ov_ppp") if _runtime == "openvino" else None | |
| if ov_ppp is not None and input_channels != 1: | |
| input_arr, scale_val, pad = _preprocess_ov_uint8(img, ov_ppp) | |
| input_dict[input_name] = input_arr | |
| result = (input_dict, scale_val, pad) | |
| if bench_fid is not None: | |
| _bench_preprocess_cache[(bench_fid, input_size)] = result | |
| return result | |
| # uint8-input ONNX models: super-gradients / Frigate YOLO-NAS exports | |
| # bake the /255 preprocessing INTO the graph and declare a uint8 input | |
| # tensor. Feeding the historical float32 [0,1] tensor is a hard | |
| # onnxruntime type error, so feed raw uint8 pixels in the declared | |
| # layout instead (`_input_dtype` is captured at session load). | |
| if config.get("_input_dtype") == "uint8" and input_channels == 3: | |
| in_side = max(input_w, input_h) | |
| if preprocess_mode == "letterbox": | |
| canvas, scale_val, pad = letterbox_image(img, in_side) | |
| else: | |
| canvas = ( | |
| img | |
| if img.size == (input_w, input_h) | |
| else img.resize((input_w, input_h), Image.BILINEAR) | |
| ) | |
| scale_val, pad = 1.0, (0, 0) | |
| pix = np.asarray(canvas, dtype=np.uint8) | |
| if pix.ndim == 2: | |
| pix = np.stack([pix] * 3, axis=-1) | |
| if _is_channels_last(config.get("_input_shape")): | |
| input_dict[input_name] = pix[np.newaxis] | |
| else: | |
| input_dict[input_name] = pix.transpose(2, 0, 1)[np.newaxis] | |
| result = (input_dict, scale_val, pad) | |
| if bench_fid is not None: | |
| _bench_preprocess_cache[(bench_fid, input_size)] = result | |
| return result | |
| if input_channels == 1: | |
| # Grayscale CTC recognizer (EasyOCR plate-OCR): aspect-preserving | |
| # resize + edge-pad + [-1,1] norm → [1,1,H,W]. NOT the squared | |
| # 3-channel RGB tensor (which the runtime rejected — the reason | |
| # plate recognition never ran). | |
| input_dict[input_name] = _preprocess_ctc_gray(img, input_w, input_h) | |
| scale_val = 1.0 | |
| pad = (0, 0) | |
| else: | |
| if preprocess_mode == "letterbox": | |
| arr, scale_val, pad = letterbox(img, input_size) | |
| else: | |
| arr = resize_image(img, input_w, input_h) | |
| scale_val = 1.0 | |
| pad = (0, 0) | |
| # Per-model input normalization (imagenet mean/std for the animal + | |
| # vehicle classifiers; no-op for detectors, CLIP, bird, etc.). | |
| arr = _apply_normalization(arr, config) | |
| # Feed the layout the model actually declares. NHWC models (ArcFace | |
| # face-embedding [N,112,112,3], CLIP) must NOT be NCHW-transposed — | |
| # a strict backend (NPU) throws on the wrong shape and the embedding | |
| # is never produced (the multi-day face-collection outage). Unknown | |
| # shape → NCHW (the prior behaviour), so detectors are unaffected. | |
| input_dict[input_name] = _to_model_tensor( | |
| arr, _is_channels_last(config.get("_input_shape")) | |
| ) | |
| result = (input_dict, scale_val, pad) | |
| if bench_fid is not None: | |
| _bench_preprocess_cache[(bench_fid, input_size)] = result | |
| return result | |
| def _postprocess( | |
| predictions: dict, | |
| slot_config: dict, | |
| orig_w: int, | |
| orig_h: int, | |
| scale_val: float, | |
| pad: tuple[int, int], | |
| elapsed_ms: float, | |
| preprocess_ms: float = 0.0, | |
| predict_ms: float = 0.0, | |
| batch_size: int = 1, | |
| ) -> dict: | |
| postprocessor_type = slot_config.get("postprocessor", "yolo") | |
| postprocessor_fn = POSTPROCESSORS.get(postprocessor_type) | |
| if postprocessor_fn is None: | |
| raise ValueError(f"Unknown postprocessor: {postprocessor_type}") | |
| result = postprocessor_fn(predictions, slot_config, orig_w, orig_h, scale_val, pad) | |
| result["inferenceMs"] = round(elapsed_ms, 2) | |
| result["preprocessMs"] = round(preprocess_ms, 2) | |
| result["predictMs"] = round(predict_ms, 2) | |
| result["batchSize"] = batch_size | |
| result["frameSize"] = f"{orig_w}x{orig_h}" | |
| return result | |
| # --------------------------------------------------------------------------- | |
| # Runtime dispatcher — pipelined prepare → predict → postprocess | |
| # --------------------------------------------------------------------------- | |
| class RuntimeDispatcher: | |
| """Split-stage dispatcher that lets preprocess run ahead of predict. | |
| Stages use separate thread pools so an in-flight `predict` doesn't | |
| block the preprocess for the next frame: | |
| camera A: preprocess → predict ─┐ | |
| camera B: preprocess → predict ─┐ | |
| camera C: preprocess → predict ─┐ | |
| While predict for A runs on the predict_pool, preprocess for B can | |
| run on the prepare_pool; each stage releases the Python event loop | |
| between calls so the reader stays responsive. | |
| `concurrency` sets the prepare pool size AND the predict pool size. | |
| For CoreML/OpenVINO the runtime itself serialises access to the | |
| device, so a single predict worker is usually enough. ONNX Runtime | |
| benefits from >1 (its InferenceSession.run is thread-safe and the | |
| native thread pool schedules concurrent runs). | |
| """ | |
| def __init__(self, runtime: str, concurrency: int) -> None: | |
| workers = max(1, concurrency) | |
| self._workers = workers | |
| self._prepare_pool = concurrent.futures.ThreadPoolExecutor( | |
| max_workers=workers, | |
| thread_name_prefix=f"{runtime}-prep", | |
| ) | |
| self._predict_pool = concurrent.futures.ThreadPoolExecutor( | |
| max_workers=workers, | |
| thread_name_prefix=f"{runtime}-predict", | |
| ) | |
| def workers(self) -> int: | |
| return self._workers | |
| async def run(self, slot: ModelSlot, img: Image.Image) -> dict: | |
| loop = asyncio.get_event_loop() | |
| t_start = time.perf_counter() | |
| orig_w, orig_h = img.size | |
| # Stage 1 — preprocess. Runs on prepare_pool; overlaps with | |
| # any predict already scheduled for a previous frame. | |
| input_dict, scale_val, pad = await loop.run_in_executor( | |
| self._prepare_pool, _preprocess, img, slot.config, | |
| ) | |
| t_predict_start = time.perf_counter() | |
| preprocess_ms = (t_predict_start - t_start) * 1000 | |
| # Stage 2 — predict. Dedicated pool so preprocessors for later | |
| # frames can run while the runtime works on this one. | |
| predictions = await loop.run_in_executor( | |
| self._predict_pool, slot.predict_fn, input_dict, | |
| ) | |
| # Stage 3 — postprocess. Cheap CPU-bound step; reuse the prepare | |
| # pool to avoid tying up a predict worker on scalar math. | |
| t_predict_end = time.perf_counter() | |
| predict_ms = (t_predict_end - t_predict_start) * 1000 | |
| elapsed_ms = (t_predict_end - t_start) * 1000 | |
| return await loop.run_in_executor( | |
| self._prepare_pool, | |
| _postprocess, | |
| predictions, slot.config, orig_w, orig_h, scale_val, pad, elapsed_ms, | |
| preprocess_ms, predict_ms, 1, | |
| ) | |
| async def run_cached( | |
| self, | |
| slot: ModelSlot, | |
| img: Image.Image, | |
| cache_key: tuple[int, int], | |
| prep_cache: dict[tuple[int, int], tuple[dict, float, tuple[int, int], int, int]], | |
| ) -> dict: | |
| """Like run() but caches the preprocessed input dict. Second+ | |
| calls for the same (frame_id, model_idx) skip preprocess entirely | |
| — goes straight to predict + postprocess. Saves ~10ms/call.""" | |
| loop = asyncio.get_event_loop() | |
| t_start = time.perf_counter() | |
| orig_w, orig_h = img.size | |
| cached = prep_cache.get(cache_key) | |
| if cached is not None: | |
| input_dict, scale_val, pad, _, _ = cached | |
| preprocess_ms = 0.0 | |
| else: | |
| input_dict, scale_val, pad = await loop.run_in_executor( | |
| self._prepare_pool, _preprocess, img, slot.config, | |
| ) | |
| preprocess_ms = (time.perf_counter() - t_start) * 1000 | |
| prep_cache[cache_key] = (input_dict, scale_val, pad, orig_w, orig_h) | |
| t_predict_start = time.perf_counter() | |
| predictions = await loop.run_in_executor( | |
| self._predict_pool, slot.predict_fn, input_dict, | |
| ) | |
| t_predict_end = time.perf_counter() | |
| predict_ms = (t_predict_end - t_predict_start) * 1000 | |
| elapsed_ms = (t_predict_end - t_start) * 1000 | |
| return await loop.run_in_executor( | |
| self._prepare_pool, | |
| _postprocess, | |
| predictions, slot.config, orig_w, orig_h, scale_val, pad, elapsed_ms, | |
| preprocess_ms, predict_ms, 1, | |
| ) | |
| async def run_list(self, slot: ModelSlot, imgs: list[Image.Image]) -> list[dict]: | |
| """Batch predict path — single CoreML/OV/ORT predict call processing | |
| all N items, then split + per-item postprocess. | |
| For CoreML: passes a list of input dicts to `model.predict(...)`, | |
| which CoreML iterates internally with shared session state — saves | |
| per-call ANE setup overhead vs a thread-pool fan-out. | |
| For OpenVINO/ONNX: falls back to a loop on the same predict thread | |
| (no model-level batching since their predict_fn signatures take one | |
| array). Still saves the asyncio.gather + ThreadPoolExecutor dispatch | |
| overhead. | |
| """ | |
| if not imgs: | |
| return [] | |
| loop = asyncio.get_event_loop() | |
| t_start = time.perf_counter() | |
| orig_sizes = [img.size for img in imgs] | |
| batch_size = len(imgs) | |
| # Stage 1 — preprocess all in parallel on the prepare pool. | |
| prepared = await asyncio.gather(*[ | |
| loop.run_in_executor(self._prepare_pool, _preprocess, img, slot.config) | |
| for img in imgs | |
| ]) | |
| # prepared = list of (input_dict, scale_val, pad) | |
| t_predict_start = time.perf_counter() | |
| preprocess_ms = (t_predict_start - t_start) * 1000 | |
| # Stage 2 — single predict call on the predict thread. | |
| runtime = _runtime | |
| supports_batch = bool(slot.config.get("_supports_batch")) | |
| max_batch = int(slot.config.get("_max_batch", 1)) | |
| if runtime == "coreml" and supports_batch and batch_size <= max_batch: | |
| # FAST PATH — stack the N preprocessed (1,3,H,W) tensors | |
| # into a single (N,3,H,W) tensor and issue ONE predict call. | |
| # The model was exported with a flexible batch axis so this | |
| # produces a single ANE dispatch processing all N frames in | |
| # parallel — amortises CoreML session setup AND the per-call | |
| # compute_units context-switch. | |
| input_name = slot.config.get("_input_name", "image") | |
| stacked = np.concatenate([p[0][input_name] for p in prepared], axis=0) | |
| def stacked_predict() -> list[dict]: | |
| names = slot.input_names | |
| payload: dict = {input_name: stacked} | |
| # Some models accept side-channel inputs (iouThreshold | |
| # etc.). Forward only those that the current model | |
| # accepts; the stacked input is set above. | |
| first = prepared[0][0] | |
| for k, v in first.items(): | |
| if k != input_name and (not names or k in names): | |
| payload[k] = v | |
| out = slot.model.predict(payload) | |
| # Split each output along axis 0 back into N items. | |
| # Output shape (N, ...); slice per index and rebuild | |
| # per-item dicts so postprocess sees the same shape it | |
| # would for a single-frame predict. | |
| results: list[dict] = [] | |
| for i in range(batch_size): | |
| item: dict = {} | |
| for k, v in out.items(): | |
| if hasattr(v, "shape") and len(v.shape) > 0 and v.shape[0] == batch_size: | |
| item[k] = v[i:i + 1] | |
| else: | |
| item[k] = v | |
| results.append(item) | |
| return results | |
| predictions_list = await loop.run_in_executor(self._predict_pool, stacked_predict) | |
| elif runtime == "coreml": | |
| # Fallback — model.predict([list of input dicts]) iterates | |
| # internally on shared session state. Not as fast as stacked | |
| # but still saves the asyncio.gather + ThreadPoolExecutor | |
| # dispatch overhead vs N separate calls. | |
| def batched_predict() -> list[dict]: | |
| names = slot.input_names | |
| if names: | |
| cleaned = [{k: v for k, v in p[0].items() if k in names} for p in prepared] | |
| else: | |
| cleaned = [p[0] for p in prepared] | |
| return slot.model.predict(cleaned) | |
| predictions_list = await loop.run_in_executor(self._predict_pool, batched_predict) | |
| else: | |
| # OpenVINO / ONNX — loop on the predict thread to keep IPC | |
| # amortisation (one executor dispatch instead of N) without | |
| # requiring model-level batched inputs. | |
| def looped_predict() -> list[dict]: | |
| return [slot.predict_fn(p[0]) for p in prepared] | |
| predictions_list = await loop.run_in_executor(self._predict_pool, looped_predict) | |
| # Stage 3 — postprocess each result. Per-item elapsed shares the | |
| # batch wall (fair attribution) — callers see it as `inferenceMs`. | |
| t_predict_end = time.perf_counter() | |
| predict_ms = (t_predict_end - t_predict_start) * 1000 | |
| elapsed_ms = (t_predict_end - t_start) * 1000 | |
| return await asyncio.gather(*[ | |
| loop.run_in_executor( | |
| self._prepare_pool, | |
| _postprocess, | |
| predictions, slot.config, | |
| orig_sizes[i][0], orig_sizes[i][1], | |
| prepared[i][1], prepared[i][2], | |
| elapsed_ms, | |
| preprocess_ms, predict_ms, batch_size, | |
| ) | |
| for i, predictions in enumerate(predictions_list) | |
| ]) | |
| def close(self) -> None: | |
| self._prepare_pool.shutdown(wait=False) | |
| self._predict_pool.shutdown(wait=False) | |
| # --------------------------------------------------------------------------- | |
| # Per-model backpressure — bound in-flight inference, drop-oldest shedding | |
| # --------------------------------------------------------------------------- | |
| # Frames a single model may hold WAITING behind its running set before the | |
| # OLDEST waiting frame is shed. Running slots (= predict-pool size) keep the | |
| # device fed; a couple of pending frames cover dispatch gaps. Anything beyond | |
| # that only sits in an invisible executor queue inflating per-frame | |
| # "inference" time — on an overloaded node (N100 iGPU) that queue wait | |
| # ballooned to 9-22s per frame. Shedding it instead makes overload read as | |
| # `dropped` frames (fast `"dropped": true` responses) rather than latency. | |
| MAX_PENDING_PER_MODEL: int = 2 | |
| class _ModelQueueState: | |
| """Mutable per-model counters for ModelBackpressure.""" | |
| running: int = 0 | |
| pending: "deque[Any]" = field(default_factory=deque) | |
| class ModelBackpressure: | |
| """Per-model in-flight bound with drop-oldest shedding. | |
| Pure decision core — no asyncio, no I/O — so it is unit-testable | |
| (see test_inference_pool_backpressure.py). The event loop calls: | |
| admit(model_idx, item) on every single-frame inference arrival. | |
| Returns (to_run, to_drop): items to dispatch NOW and items shed | |
| (oldest queued first). At most one of the lists is non-empty. | |
| complete(model_idx) when a dispatched inference finishes. | |
| Returns the next queued items to dispatch (0 or 1). | |
| Why per-model and why these bounds: the predict pool has `workers` | |
| threads shared by all models; a model gets at most `max_running` | |
| concurrently-dispatched inferences (sized to the predict pool), plus | |
| a `max_pending` waiting buffer. The previous GLOBAL ceiling | |
| (`workers * 4`, min 16) was unreachable from the live pipeline — the | |
| TS runner's semaphore keeps at most `maxConcurrentInferences` frames | |
| in flight per node, well below 16 — so overload never shed frames; it | |
| just queued inside the executor as multi-second "inference" time. | |
| """ | |
| def __init__(self, max_running: int, max_pending: int) -> None: | |
| self._max_running = max(1, max_running) | |
| self._max_pending = max(0, max_pending) | |
| self._states: dict[int, _ModelQueueState] = {} | |
| def _state(self, model_idx: int) -> _ModelQueueState: | |
| state = self._states.get(model_idx) | |
| if state is None: | |
| state = _ModelQueueState() | |
| self._states[model_idx] = state | |
| return state | |
| def admit(self, model_idx: int, item: Any) -> tuple[list, list]: | |
| """Admit one inference request; returns (to_run, to_drop).""" | |
| state = self._state(model_idx) | |
| if state.running < self._max_running: | |
| state.running += 1 | |
| return [item], [] | |
| state.pending.append(item) | |
| dropped: list = [] | |
| while len(state.pending) > self._max_pending: | |
| # DROP-OLDEST: the newest frame is the most valuable one for a | |
| # live detection pipeline — shed the stalest queued frame(s). | |
| dropped.append(state.pending.popleft()) | |
| return [], dropped | |
| def complete(self, model_idx: int) -> list: | |
| """One running inference finished; returns queued items to start.""" | |
| state = self._state(model_idx) | |
| state.running = max(0, state.running - 1) | |
| if state.pending and state.running < self._max_running: | |
| state.running += 1 | |
| return [state.pending.popleft()] | |
| return [] | |
| def snapshot(self) -> dict: | |
| """Per-model {running, pending} — queue depth for the mem_stats line. | |
| A leak shows RSS climbing while these stay flat; a genuinely growing | |
| working set shows them climbing together. Pure read, no mutation. | |
| """ | |
| return { | |
| str(idx): {"running": st.running, "pending": len(st.pending)} | |
| for idx, st in self._states.items() | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Memory diagnostics (mem_stats command) | |
| # | |
| # The pools leak over hours and get OOM-killed (2026-08-18: anon-rss 3.8 GB / | |
| # 8.8 GB, twice in six hours) — the Node-side watchdog samples RSS from | |
| # /proc/<pid>/status, and asks THIS process once per sweep for the internals | |
| # /proc cannot see: what the pool has been asked to do (request counters), | |
| # what it is holding (caches, queue depths) and what the allocator is doing | |
| # (arena count — the lead on total-vm 47.5 GB vs 11.3 GB at similar RSS). | |
| # Everything here is a cheap O(1)-ish read; nothing walks the heap. | |
| # --------------------------------------------------------------------------- | |
| class PoolCounters: | |
| """Requests served since spawn, by opcode. Monotonic, wraps never.""" | |
| infer_jpeg: int = 0 | |
| infer_raw: int = 0 | |
| infer_batch: int = 0 | |
| infer_cached: int = 0 | |
| commands: int = 0 | |
| frames_cached: int = 0 | |
| # Requests answered `dropped` because their wire deadline had already | |
| # passed when they would have run — dead work NOT executed (D350). For a | |
| # batch, every item counts. | |
| shed_deadline_expired: int = 0 | |
| def as_dict(self) -> dict: | |
| return { | |
| "inferJpeg": self.infer_jpeg, | |
| "inferRaw": self.infer_raw, | |
| "inferBatch": self.infer_batch, | |
| "inferCached": self.infer_cached, | |
| "commands": self.commands, | |
| "framesCached": self.frames_cached, | |
| "shedDeadlineExpired": self.shed_deadline_expired, | |
| "totalInference": ( | |
| self.infer_jpeg + self.infer_raw + self.infer_batch + self.infer_cached | |
| ), | |
| } | |
| def _read_self_memory() -> dict: | |
| """RSS / VmSize / peak-RSS / swap / thread count for THIS process. | |
| /proc on Linux (one 4 KB read); getrusage maxrss everywhere as backstop. | |
| Partial results are fine — an empty dict means the read failed entirely. | |
| """ | |
| out: dict = {} | |
| try: | |
| import resource | |
| ru = resource.getrusage(resource.RUSAGE_SELF) | |
| # ru_maxrss is kilobytes on Linux, BYTES on macOS. | |
| out["maxRssKb"] = ( | |
| int(ru.ru_maxrss // 1024) if sys.platform == "darwin" else int(ru.ru_maxrss) | |
| ) | |
| except Exception: | |
| pass | |
| try: | |
| with open("/proc/self/status", "r") as fh: | |
| text = fh.read() | |
| for label, key in ( | |
| ("VmRSS", "rssKb"), | |
| ("VmSize", "vmKb"), | |
| ("VmHWM", "hwmKb"), | |
| ("VmSwap", "swapKb"), | |
| ): | |
| match = re.search(rf"^{label}:\s+(\d+)\s*kB", text, re.M) | |
| if match: | |
| out[key] = int(match.group(1)) | |
| match = re.search(r"^Threads:\s+(\d+)", text, re.M) | |
| if match: | |
| out["threads"] = int(match.group(1)) | |
| except Exception: | |
| pass | |
| return out | |
| def _malloc_arena_count() -> Optional[int]: | |
| """Count glibc malloc arenas via malloc_info(3) into an in-memory stream. | |
| Each arena reserves up to 64 MB of VIRTUAL space — the count is the direct | |
| test of the "total-vm 47.5 GB is allocator arenas" hypothesis from the | |
| 2026-08-18 OOM audit. glibc/Linux only; None elsewhere or on any failure | |
| (never raises — this rides the mem_stats reply and must not kill it). | |
| """ | |
| if not sys.platform.startswith("linux"): | |
| return None | |
| try: | |
| import ctypes | |
| libc = ctypes.CDLL(None, use_errno=True) | |
| libc.open_memstream.restype = ctypes.c_void_p | |
| libc.open_memstream.argtypes = [ctypes.c_void_p, ctypes.c_void_p] | |
| libc.malloc_info.argtypes = [ctypes.c_int, ctypes.c_void_p] | |
| libc.fclose.argtypes = [ctypes.c_void_p] | |
| libc.free.argtypes = [ctypes.c_void_p] | |
| buf = ctypes.c_char_p() | |
| size = ctypes.c_size_t() | |
| stream = libc.open_memstream(ctypes.byref(buf), ctypes.byref(size)) | |
| if not stream: | |
| return None | |
| rc = libc.malloc_info(0, stream) | |
| libc.fclose(stream) # flushes and finalises buf/size | |
| try: | |
| if rc != 0 or not buf.value: | |
| return None | |
| xml = ctypes.string_at(buf, size.value) | |
| return xml.count(b"<heap nr=") | |
| finally: | |
| libc.free(buf) | |
| except Exception: | |
| return None | |
| def _mem_stats_payload( | |
| models: list, | |
| counters: PoolCounters, | |
| backpressure: "ModelBackpressure", | |
| frame_cache_len: int, | |
| preprocess_cache_len: int, | |
| bench_cache_len: int, | |
| uptime_sec: int, | |
| model_load: Optional[dict] = None, | |
| ) -> dict: | |
| """Assemble the mem_stats reply. Pure given its inputs except the two | |
| process-wide reads (memory + arenas), so the shape is unit-testable.""" | |
| return { | |
| "cmd": "mem_stats", | |
| "status": "ok", | |
| "memory": _read_self_memory(), | |
| "mallocArenas": _malloc_arena_count(), | |
| "gcCounts": list(gc.get_count()), | |
| "counters": counters.as_dict(), | |
| "backpressure": backpressure.snapshot(), | |
| "modelsLoaded": sum(1 for s in models if getattr(s, "loaded", False)), | |
| "frameCacheEntries": frame_cache_len, | |
| "preprocessCacheEntries": preprocess_cache_len, | |
| "benchPreprocessCacheEntries": bench_cache_len, | |
| "uptimeSec": uptime_sec, | |
| # What is compiling, and whether a compile ever failed to return | |
| # (D653) — the one fact a host needs to tell "slow" from "poisoned". | |
| "modelLoad": model_load, | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Single-frame inference execution — module-level so the deadline gate is | |
| # unit-testable (test_inference_pool_deadline.py) without driving the whole | |
| # event loop. `handle_inference` in _run() delegates here; the caller's | |
| # `finally` still owns the backpressure slot release. | |
| # --------------------------------------------------------------------------- | |
| async def _execute_inference( | |
| req_id: int, | |
| img: Any, | |
| model_idx: int, | |
| deadline_ms: int, | |
| *, | |
| models: "list[ModelSlot]", | |
| dispatcher: Any, | |
| writer: Any, | |
| batch_mode: str, | |
| get_accumulator: "Optional[Callable[[int], Any]]", | |
| counters: "PoolCounters", | |
| ) -> None: | |
| """Deadline gate + predict for ONE admitted single-frame request. | |
| The deadline is checked HERE — immediately before predict — rather than | |
| only at arrival, because this coroutine is entered both on admission and | |
| when the backpressure queue promotes a waiting frame: the wait is exactly | |
| where a request goes stale. The gate comes first (before the model-loaded | |
| check) on purpose — the cheapest question first, and an abandoned request | |
| deserves a shed, not a diagnosis. `dispatcher` and `writer` are | |
| duck-typed (RuntimeDispatcher / ResponseWriter in production). | |
| """ | |
| now_ms = time.time() * 1000.0 | |
| if _deadline_expired(deadline_ms, now_ms): | |
| counters.shed_deadline_expired += 1 | |
| _log_deadline_shed(counters.shed_deadline_expired, model_idx, now_ms - deadline_ms) | |
| await writer.send(req_id, _dropped_deadline_payload()) | |
| return | |
| if model_idx >= len(models) or not models[model_idx].loaded: | |
| await writer.send(req_id, { | |
| "error": f"Model {model_idx} not loaded", | |
| "kind": "detections", | |
| "detections": [], | |
| "inferenceMs": 0, | |
| }) | |
| return | |
| if batch_mode == "window" and get_accumulator is not None: | |
| await get_accumulator(model_idx).submit(req_id, img) | |
| return | |
| result = await dispatcher.run(models[model_idx], img) | |
| await writer.send(req_id, result) | |
| # --------------------------------------------------------------------------- | |
| # IPC — binary framing with request_id multiplexing | |
| # --------------------------------------------------------------------------- | |
| HEADER_LEN = 4 # total_len prefix | |
| PREFIX_LEN = HEADER_LEN + 4 + 1 # total_len + req_id + msg_type | |
| async def _read_exact(reader: asyncio.StreamReader, n: int) -> bytes: | |
| return await reader.readexactly(n) | |
| class ResponseWriter: | |
| """Serialises writes to the stdout pipe; safe from multiple coroutines.""" | |
| def __init__(self, writer: asyncio.StreamWriter) -> None: | |
| self._writer = writer | |
| self._lock = asyncio.Lock() | |
| async def send(self, req_id: int, payload: dict) -> None: | |
| # numpy scalars (float32 / int64 / bool_) are NOT JSON-serializable | |
| # under the stdlib encoder. With numpy 2.x the previous lenient | |
| # path is gone, so we must coerce before encoding. Cheap default | |
| # callback only kicks in for unhandled types — Python natives skip | |
| # it entirely. | |
| data = json.dumps(payload, default=_json_default).encode("utf-8") | |
| total_len = 4 + len(data) # req_id(4) + json | |
| header = struct.pack("<II", total_len, req_id) | |
| async with self._lock: | |
| self._writer.write(header) | |
| self._writer.write(data) | |
| await self._writer.drain() | |
| def _json_default(obj: Any) -> Any: | |
| """Coerce numpy scalars/arrays to JSON-friendly Python natives.""" | |
| if isinstance(obj, np.generic): | |
| return obj.item() | |
| if isinstance(obj, np.ndarray): | |
| return obj.tolist() | |
| raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable") | |
| # --------------------------------------------------------------------------- | |
| # Command handlers | |
| # --------------------------------------------------------------------------- | |
| # Config keys a LOADED slot can absorb through `reconfigure` — decode-scoped | |
| # values the postprocessor re-reads from `slot.config` on every inference, so | |
| # they take effect on the next frame with no session reload. Mirrors | |
| # `PoolDecodeSettings` on the TypeScript side; adding a key here without adding | |
| # it there (or vice versa) reintroduces the D214 dead-knob class. | |
| RECONFIGURABLE_CONFIG_KEYS = frozenset({"nmsIouThreshold"}) | |
| def _handle_command(models: list[ModelSlot], cmd: dict) -> dict: | |
| """The model commands that need no compile: unload, reconfigure, status. | |
| `load` and `replace` are NOT here — they compile, and a compile must never | |
| run on the event loop (D653). They go through `_handle_model_command`, | |
| which the `CommandRouter` awaits off the loop's reader. | |
| """ | |
| action = cmd.get("cmd") | |
| if action == "unload": | |
| index = cmd["index"] | |
| if index < len(models) and models[index].loaded: | |
| _unload_model(models[index]) | |
| sys.stderr.write(f"Model {index} unloaded\n") | |
| sys.stderr.flush() | |
| return {"cmd": "unload", "index": index, "status": "ok"} | |
| if action == "reconfigure": | |
| # Retune a LOADED slot's decode-scoped settings in place — no session | |
| # re-instantiation. `slot.config` is re-read by the postprocessor on | |
| # every inference (see `_postprocess`), so mutating these keys is | |
| # picked up by the next frame without a reload. | |
| # | |
| # ONLY the allow-listed keys are accepted. A `path` or `inputSize` | |
| # arriving here would be written into the config while the loaded | |
| # session kept serving the OLD model — a silent, undetectable lie. | |
| # Anything else is refused loudly; the caller must use `replace`. | |
| index = cmd["index"] | |
| patch = cmd.get("config") or {} | |
| unknown = sorted(k for k in patch if k not in RECONFIGURABLE_CONFIG_KEYS) | |
| if unknown: | |
| return { | |
| "cmd": "reconfigure", | |
| "index": index, | |
| "status": "error", | |
| "error": ( | |
| f"keys not reconfigurable without a reload: {', '.join(unknown)} " | |
| f"(allowed: {', '.join(sorted(RECONFIGURABLE_CONFIG_KEYS))})" | |
| ), | |
| } | |
| if index >= len(models) or not models[index].loaded: | |
| return { | |
| "cmd": "reconfigure", | |
| "index": index, | |
| "status": "error", | |
| "error": f"model {index} is not loaded", | |
| } | |
| slot = models[index] | |
| before = {k: slot.config.get(k) for k in patch} | |
| slot.config.update(patch) | |
| sys.stderr.write(f"Model {index} reconfigured: {before} -> {dict(patch)}\n") | |
| sys.stderr.flush() | |
| return {"cmd": "reconfigure", "index": index, "status": "ok"} | |
| if action == "status": | |
| status = [] | |
| for i, slot in enumerate(models): | |
| status.append({ | |
| "index": i, | |
| "path": slot.config.get("path") if slot.loaded else None, | |
| "loaded": slot.loaded, | |
| "postprocessor": slot.config.get("postprocessor") if slot.loaded else None, | |
| }) | |
| return {"cmd": "status", "models": status, "benchCache": _bench_cache_stats()} | |
| return {"cmd": action or "unknown", "status": "error", "error": f"Unknown command: {action}"} | |
| # --------------------------------------------------------------------------- | |
| # Model loads run OFF the event loop, one at a time, under a hard bound (D653) | |
| # --------------------------------------------------------------------------- | |
| # | |
| # 2026-09-26, hub `openvino:gpu`: one `core.compile_model()` of | |
| # `camstack-yunet-2023mar.xml` for the iGPU never returned. It ran INLINE on this | |
| # loop, so from 18:33:48 until a container restart at 22:41 the worker read no | |
| # request and wrote no reply: 786 client-side timeouts on six cameras, 123 of | |
| # 123 `mem_stats` probes lost, and a process that looked perfectly alive. | |
| # | |
| # So a load now runs on its own DAEMON thread and the loop only awaits it: | |
| # inference and `mem_stats` keep being served while a compile runs. And it is | |
| # BOUNDED: past `MODEL_LOAD_TIMEOUT_SEC` the load answers `compile-timeout`, | |
| # naming the model, instead of never answering. | |
| # | |
| # A native thread stuck in a GPU compiler cannot be cancelled or killed from | |
| # Python. So a timeout POISONS the worker FOR NEW LOADS: every later load is | |
| # refused at once (`worker-poisoned`) rather than queued behind a thread that | |
| # may never free the compiler, and the models already loaded keep serving. | |
| # | |
| # The abandoned compile is NOT abandoned work (fix round 1): it keeps running, | |
| # so a slow but finite compile still writes its cache. When it comes back the | |
| # worker says so on the event channel (`compile-finished-late`) and is usable | |
| # for loads again — the next load of that model is a warm cache hit. If it | |
| # never comes back, the host kills the process at its HARD bound | |
| # (shared-inference-pool.ts `POOL_MODEL_LOAD_BOUNDS`). Daemon, not an executor | |
| # thread: `concurrent.futures` joins its threads at interpreter exit, so a stuck | |
| # compile would also have made the worker impossible to shut down. | |
| # SOFT bound on ONE model load — compile included — in seconds. A cold GPU | |
| # compile of the whole model set on the N100 iGPU measured ~25 s; one model is | |
| # normally 1-3 s. The host declares its own value in the startup config | |
| # (`modelLoadTimeoutMs`) so its reply deadline and this bound cannot drift; | |
| # this is what an old host that declares nothing gets. | |
| MODEL_LOAD_TIMEOUT_SEC = 120.0 | |
| # Request id reserved for UNSOLICITED worker events (a reply to no request). | |
| # The host never allocates it (shared-inference-pool.ts `WORKER_EVENT_REQ_ID`). | |
| WORKER_EVENT_REQ_ID = 0xFFFFFFFF | |
| # Commands answered by the loop immediately, never queued behind a load: the | |
| # probe the host uses to tell "slow" from "wedged" must never wait for the very | |
| # compile it is asking about. | |
| INLINE_COMMANDS = frozenset({"mem_stats", "uncache_frame"}) | |
| # Commands that compile a model — the only ones that leave the loop. | |
| LOAD_COMMANDS = frozenset({"load", "replace"}) | |
| def _model_load_timeout_sec(config: dict) -> float: | |
| """The load bound for this connection: the host's, else the default.""" | |
| raw = config.get("modelLoadTimeoutMs") | |
| if isinstance(raw, (int, float)) and not isinstance(raw, bool) and raw > 0: | |
| return float(raw) / 1000.0 | |
| return MODEL_LOAD_TIMEOUT_SEC | |
| def _model_id_of(config: dict) -> str: | |
| """The model's file stem — what a log line and the host's reply carry.""" | |
| path = str(config.get("path") or "") | |
| stem = os.path.splitext(os.path.basename(path))[0] | |
| return stem or "unknown" | |
| class ModelLoadTimeout(Exception): | |
| """A load did not finish inside its bound; new loads are refused until it does.""" | |
| class WorkerPoisoned(Exception): | |
| """A previous load is still running past its bound; no new load starts.""" | |
| class ModelLoader: | |
| """Runs `_load_model` on a daemon thread under a hard bound. | |
| One load at a time is the CALLER's guarantee (the `CommandRouter` consumes | |
| model commands in order); this class only owns the thread, the bound and | |
| the poisoning. `load_fn` is injectable for tests; production passes | |
| nothing and gets the module's `_load_model`, looked up at call time. | |
| `notify` sends an unsolicited event to the host (the `compile-finished-late` | |
| line); absent, the event is only written to stderr. | |
| """ | |
| def __init__( | |
| self, | |
| timeout_sec: float = MODEL_LOAD_TIMEOUT_SEC, | |
| load_fn: "Optional[Callable[[ModelSlot, dict], None]]" = None, | |
| notify: "Optional[Callable[[dict], Awaitable[None]]]" = None, | |
| ) -> None: | |
| self.timeout_sec = timeout_sec | |
| self._load_fn = load_fn | |
| self._notify = notify | |
| self.finished_late = 0 | |
| # Test seam: runs on the load thread right after it has claimed the | |
| # outcome — lets a test hold the thread exactly in the race window. | |
| self._after_claim: Optional[Callable[[], None]] = None | |
| # Monotonic start of the compile running past its bound, if any — the | |
| # host counts its HARD bound from here, not from a later refusal. | |
| self._abandoned_started: Optional[float] = None | |
| # `ensure_future` tasks are only weakly referenced by the loop; keep | |
| # each notification alive until it is sent. | |
| self._pending_notifications: "set[asyncio.Future[None]]" = set() | |
| self.in_flight: Optional[dict] = None | |
| self.poisoned_by: Optional[dict] = None | |
| self.completed = 0 | |
| self.failed = 0 | |
| self.timed_out = 0 | |
| def snapshot(self) -> dict: | |
| """What the `mem_stats` reply carries about loads.""" | |
| in_flight = None | |
| if self.in_flight is not None: | |
| in_flight = { | |
| "index": self.in_flight["index"], | |
| "modelId": self.in_flight["modelId"], | |
| "elapsedMs": round((time.monotonic() - self.in_flight["startedAt"]) * 1000), | |
| } | |
| return { | |
| "inFlight": in_flight, | |
| "poisonedBy": dict(self.poisoned_by) if self.poisoned_by is not None else None, | |
| "boundMs": round(self.timeout_sec * 1000), | |
| "completed": self.completed, | |
| "failed": self.failed, | |
| "timedOut": self.timed_out, | |
| "finishedLate": self.finished_late, | |
| } | |
| def _finish_abandoned( | |
| self, index: int, model_id: str, outcome: Optional[BaseException], elapsed_ms: int, | |
| ) -> None: | |
| """On the loop: an abandoned compile came back. The compiler is free | |
| again, so the worker loads again; a successful one has written its | |
| cache, so the host's next load of the model is warm. Its RESULT is | |
| still discarded — the host was told it failed and has moved on.""" | |
| self.finished_late += 1 | |
| self.poisoned_by = None | |
| self._abandoned_started = None | |
| payload = { | |
| "event": "compile-finished-late", | |
| "index": index, | |
| "modelId": model_id, | |
| "ok": outcome is None, | |
| "elapsedMs": elapsed_ms, | |
| } | |
| if outcome is not None: | |
| payload["error"] = str(outcome) | |
| if self._notify is not None: | |
| task = asyncio.ensure_future(self._notify(payload)) | |
| self._pending_notifications.add(task) | |
| task.add_done_callback(self._pending_notifications.discard) | |
| def abandoned_elapsed_ms(self) -> Optional[int]: | |
| """How long the compile running past its bound has been running.""" | |
| if self._abandoned_started is None: | |
| return None | |
| return round((time.monotonic() - self._abandoned_started) * 1000) | |
| async def load(self, slot: ModelSlot, config: dict, index: int) -> None: | |
| """Load `config` into the DETACHED `slot`, or raise. | |
| Raises `WorkerPoisoned` without starting anything once a previous load | |
| timed out, `ModelLoadTimeout` when this one exceeds the bound, and the | |
| load's own exception when it fails in time. | |
| """ | |
| model_id = _model_id_of(config) | |
| if self.poisoned_by is not None: | |
| raise WorkerPoisoned( | |
| f"worker poisoned for loads: the load of {self.poisoned_by['modelId']} is still " | |
| f"running past its bound ({self.poisoned_by['boundMs']}ms) — no new load starts " | |
| f"until it returns" | |
| ) | |
| load_fn = self._load_fn if self._load_fn is not None else _load_model | |
| loop = asyncio.get_running_loop() | |
| done: "asyncio.Future[None]" = loop.create_future() | |
| # Exactly ONE of "the load returned in time" and "the bound fired" | |
| # wins, decided under this lock (release review). Without it the | |
| # thread could see "not abandoned", the bound fire before its settle | |
| # ran, and the worker stay poisoned forever for a load that finished. | |
| decision = threading.Lock() | |
| claimed: dict = {"by": None} # "thread" | "bound" | |
| def settle(exc: Optional[BaseException]) -> None: | |
| if done.done(): | |
| return | |
| if exc is None: | |
| done.set_result(None) | |
| else: | |
| done.set_exception(exc) | |
| def target() -> None: | |
| outcome: Optional[BaseException] = None | |
| try: | |
| load_fn(slot, config) | |
| except BaseException as exc: # noqa: BLE001 — handed to the awaiting loop | |
| outcome = exc | |
| with decision: | |
| late = claimed["by"] == "bound" | |
| if not late: | |
| claimed["by"] = "thread" | |
| if self._after_claim is not None: | |
| self._after_claim() | |
| if late: | |
| # The bound already fired. Say that the native call DID come | |
| # back, and when — the only evidence of how long a slow | |
| # compile really takes — and hand the worker back its loads. | |
| elapsed_ms = round((time.monotonic() - started) * 1000) | |
| sys.stderr.write( | |
| f"Model {index} load of {model_id} returned AFTER its bound " | |
| f"in {elapsed_ms}ms ({'error: ' + str(outcome) if outcome else 'ok'}) " | |
| f"— result discarded, loads re-enabled\n" | |
| ) | |
| sys.stderr.flush() | |
| try: | |
| loop.call_soon_threadsafe( | |
| self._finish_abandoned, index, model_id, outcome, elapsed_ms, | |
| ) | |
| except RuntimeError: | |
| pass # loop closed: the process is shutting down | |
| return | |
| try: | |
| loop.call_soon_threadsafe(settle, outcome) | |
| except RuntimeError: | |
| pass # loop closed: the process is shutting down | |
| started = time.monotonic() | |
| self.in_flight = {"index": index, "modelId": model_id, "startedAt": started} | |
| thread = threading.Thread(target=target, name=f"model-load-{index}", daemon=True) | |
| thread.start() | |
| try: | |
| try: | |
| await asyncio.wait_for(asyncio.shield(done), timeout=self.timeout_sec) | |
| except asyncio.TimeoutError: | |
| with decision: | |
| bound_won = claimed["by"] is None | |
| if bound_won: | |
| claimed["by"] = "bound" | |
| if bound_won: | |
| raise | |
| # The load returned AT the bound: its settle is already on | |
| # its way to this loop. It counts as in time. | |
| await done | |
| except asyncio.TimeoutError: | |
| self.timed_out += 1 | |
| self._abandoned_started = started | |
| self.poisoned_by = { | |
| "index": index, | |
| "modelId": model_id, | |
| "boundMs": round(self.timeout_sec * 1000), | |
| "at": round(time.time() * 1000), | |
| } | |
| raise ModelLoadTimeout( | |
| f"load of {model_id} (index {index}) did not return within " | |
| f"{round(self.timeout_sec * 1000)}ms — it keeps running so its cache " | |
| f"can be written; no new load starts until it returns" | |
| ) from None | |
| except Exception: | |
| self.failed += 1 | |
| raise | |
| finally: | |
| self.in_flight = None | |
| self.completed += 1 | |
| async def _handle_model_command( | |
| models: "list[ModelSlot]", cmd: dict, loader: ModelLoader, | |
| ) -> dict: | |
| """Every model-management command, in order; `load`/`replace` off the loop. | |
| The new model is loaded into a DETACHED slot and installed only once its | |
| load completed inside the bound, so a compile that returns late — after the | |
| host was told `compile-timeout` — can never surface as a model the host | |
| believes failed. Whatever the index held before keeps serving until that | |
| swap: a `replace` that fails or times out leaves the old model in place, | |
| and the old slot object is only dropped, never mutated, so a predict | |
| already holding it finishes normally. | |
| """ | |
| action = cmd.get("cmd") | |
| if action not in LOAD_COMMANDS: | |
| return _handle_command(models, cmd) | |
| index = cmd["index"] | |
| config = cmd["config"] | |
| model_id = _model_id_of(config) | |
| while len(models) <= index: | |
| models.append(ModelSlot()) | |
| slot = ModelSlot() | |
| sys.stderr.write( | |
| f"Model {index} load started: {config['path']} " | |
| f"(bound {round(loader.timeout_sec * 1000)}ms)\n" | |
| ) | |
| sys.stderr.flush() | |
| t0 = time.perf_counter() | |
| try: | |
| await loader.load(slot, config, index) | |
| except WorkerPoisoned as exc: | |
| return { | |
| "cmd": action, "index": index, "status": "error", | |
| "reason": "worker-poisoned", "modelId": model_id, "poisoned": True, | |
| # How long the EARLIER compile has run: the host's hard bound | |
| # counts from its start, not from this refusal. | |
| "compileElapsedMs": loader.abandoned_elapsed_ms(), | |
| "error": str(exc), | |
| } | |
| except ModelLoadTimeout as exc: | |
| sys.stderr.write(f"Model {index} load TIMED OUT: {exc}\n") | |
| sys.stderr.flush() | |
| return { | |
| "cmd": action, "index": index, "status": "error", | |
| "reason": "compile-timeout", "modelId": model_id, "poisoned": True, | |
| "timeoutMs": round(loader.timeout_sec * 1000), | |
| "compileElapsedMs": loader.abandoned_elapsed_ms(), | |
| "error": str(exc), | |
| } | |
| except Exception as exc: | |
| reply = {"cmd": action, "index": index, "status": "error", "error": str(exc)} | |
| # A typed refusal (e.g. `native-load-crashed`) names its class, so the | |
| # host can tell it from an ordinary failed load. | |
| reason = getattr(exc, "reason", None) | |
| if isinstance(reason, str): | |
| reply["reason"] = reason | |
| reply["modelId"] = model_id | |
| return reply | |
| models[index] = slot | |
| load_ms = round((time.perf_counter() - t0) * 1000) | |
| verb = "loaded" if action == "load" else "replaced" | |
| sys.stderr.write(f"Model {index} {verb}: {config['path']} ({load_ms}ms)\n") | |
| sys.stderr.flush() | |
| return {"cmd": action, "index": index, "status": "ok", "loadMs": load_ms} | |
| class CommandRouter: | |
| """Routes MSG_COMMAND frames without ever blocking the reader. | |
| `INLINE_COMMANDS` are answered at once by `inline` (it needs the loop's own | |
| state: caches, backpressure, counters). Everything else is queued and | |
| consumed IN ORDER by one task — model slots are indices shared with the | |
| host, so a `status` or `unload` must never overtake the `load` sent before | |
| it. | |
| """ | |
| def __init__( | |
| self, | |
| models: "list[ModelSlot]", | |
| loader: ModelLoader, | |
| send: "Callable[[int, dict], Awaitable[None]]", | |
| inline: "Callable[[dict], dict]", | |
| ) -> None: | |
| self._models = models | |
| self._loader = loader | |
| self._send = send | |
| self._inline = inline | |
| self._queue: "asyncio.Queue[tuple[int, dict]]" = asyncio.Queue() | |
| self._consumer: "Optional[asyncio.Task[None]]" = None | |
| async def submit(self, req_id: int, cmd: dict) -> None: | |
| if cmd.get("cmd") in INLINE_COMMANDS: | |
| try: | |
| response = self._inline(cmd) | |
| except Exception as exc: | |
| response = {"cmd": cmd.get("cmd") or "unknown", "status": "error", "error": str(exc)} | |
| await self._send(req_id, response) | |
| return | |
| self._queue.put_nowait((req_id, cmd)) | |
| if self._consumer is None or self._consumer.done(): | |
| self._consumer = asyncio.create_task(self._consume()) | |
| async def _consume(self) -> None: | |
| while not self._queue.empty(): | |
| req_id, cmd = self._queue.get_nowait() | |
| try: | |
| response = await _handle_model_command(self._models, cmd, self._loader) | |
| except Exception as exc: | |
| response = {"cmd": cmd.get("cmd") or "unknown", "status": "error", "error": str(exc)} | |
| await self._send(req_id, response) | |
| # --------------------------------------------------------------------------- | |
| # Main async loop | |
| # --------------------------------------------------------------------------- | |
| async def _run() -> None: | |
| loop = asyncio.get_event_loop() | |
| reader = asyncio.StreamReader(loop=loop) | |
| protocol = asyncio.StreamReaderProtocol(reader, loop=loop) | |
| await loop.connect_read_pipe(lambda: protocol, sys.stdin.buffer) | |
| transport, stdout_protocol = await loop.connect_write_pipe( | |
| asyncio.streams.FlowControlMixin, sys.stdout.buffer, | |
| ) | |
| stream_writer = asyncio.StreamWriter(transport, stdout_protocol, None, loop) | |
| writer = ResponseWriter(stream_writer) | |
| t_start = time.perf_counter() | |
| # ── Startup — first message is config (req_id=0, msg_type=COMMAND) ── | |
| header = await _read_exact(reader, PREFIX_LEN) | |
| total_len, req_id = struct.unpack("<II", header[0:8]) | |
| msg_type = header[8] | |
| remaining = total_len - 5 | |
| payload = await _read_exact(reader, remaining) | |
| if msg_type != MSG_COMMAND: | |
| sys.stderr.write(f"Expected startup config (cmd), got msg_type={msg_type}\n") | |
| sys.exit(1) | |
| config = json.loads(payload) | |
| runtime = config.get("runtime", "coreml") | |
| pool_device = str(config.get("device", "") or "") | |
| # Wire version for THIS connection (D350): an old host declares nothing | |
| # and negotiates v1 — byte-identical historical behaviour, no deadline | |
| # prefix expected on any payload. | |
| wire_version = _negotiate_wire_version(config) | |
| concurrency = int(config.get("concurrency", 1) or 1) | |
| batch_mode = str(config.get("batch_mode", "none")) | |
| if batch_mode not in ("none", "list", "window"): | |
| batch_mode = "none" | |
| window_ms = max(0, int(config.get("window_ms", 0) or 0)) | |
| max_batch_size = max(1, int(config.get("max_batch_size", 1) or 1)) | |
| sys.stderr.write( | |
| f"Initializing runtime: {runtime} (concurrency={concurrency}, " | |
| f"batch_mode={batch_mode}, window_ms={window_ms}, max_batch={max_batch_size}, " | |
| f"wire_version={wire_version})\n" | |
| ) | |
| sys.stderr.flush() | |
| _init_runtime(runtime, pool_device) | |
| # Off-loop, bounded model loads (D653). Startup-config models below still | |
| # load inline: nothing is served before `ready`, and the host bounds that | |
| # handshake itself. | |
| model_loader = ModelLoader( | |
| timeout_sec=_model_load_timeout_sec(config), | |
| notify=lambda event: writer.send(WORKER_EVENT_REQ_ID, event), | |
| ) | |
| models: list[ModelSlot] = [] | |
| for i, mc in enumerate(config.get("models", [])): | |
| slot = ModelSlot() | |
| sys.stderr.write(f"Loading model {i}: {mc['path']}\n") | |
| sys.stderr.flush() | |
| try: | |
| _load_model(slot, mc) | |
| except Exception as exc: | |
| sys.stderr.write(f"ERROR loading model {i}: {exc}\n") | |
| sys.stderr.flush() | |
| models.append(slot) | |
| # OpenVINO: size the predict pool so the fixed InferRequest pool (built per | |
| # model in _load_model) is actually fed by concurrent predict workers. The | |
| # addon passes concurrency=1 (it predates the request pool) and models load | |
| # lazily AFTER this point, so we can't read OPTIMAL_NUMBER_OF_INFER_REQUESTS | |
| # here — `_ov_request_count(0)` then answers OV_DEFAULT_CONCURRENCY, because | |
| # a device that could not be asked is UNKNOWN and not small. A startup-loaded | |
| # model that DID answer sizes it to what the device actually wants. The | |
| # per-model InferRequest pool goes through the same function, so the two | |
| # cannot drift apart. | |
| if runtime == "openvino": | |
| target = _ov_request_count(_ov_optimal_reqs) | |
| if target > concurrency: | |
| sys.stderr.write( | |
| f"OpenVINO LATENCY: predict concurrency {concurrency} -> {target}\n" | |
| ) | |
| sys.stderr.flush() | |
| concurrency = target | |
| dispatcher = RuntimeDispatcher(runtime, concurrency) | |
| # Backpressure: bound in-flight single-frame inference PER MODEL, shedding | |
| # the OLDEST queued frame when full (see ModelBackpressure). max_running is | |
| # sized to the predict pool — dispatching more than `workers` concurrent | |
| # inferences for one model can't run any faster, it only queues invisibly | |
| # inside the executor and balloons per-frame "inference" time (9-22s | |
| # observed on the N100 iGPU). Healthy nodes never fill the pending buffer, | |
| # so this is a no-op except under genuine overload. | |
| backpressure = ModelBackpressure( | |
| max_running=dispatcher.workers, | |
| max_pending=MAX_PENDING_PER_MODEL, | |
| ) | |
| # Served-request counters + spawn time for the mem_stats command — the | |
| # denominator that lets the watchdog tell "RSS grew because it worked" | |
| # from "RSS grew per unit of work" (a leak). | |
| counters = PoolCounters() | |
| pool_started_monotonic = time.monotonic() | |
| startup_ms = round((time.perf_counter() - t_start) * 1000) | |
| loaded_count = sum(1 for s in models if s.loaded) | |
| sys.stderr.write( | |
| f"Ready — {loaded_count}/{len(models)} model(s) in {startup_ms}ms " | |
| f"(runtime={runtime}, workers={dispatcher.workers})\n" | |
| ) | |
| sys.stderr.flush() | |
| await writer.send(req_id, { | |
| "status": "ready", | |
| "models": loaded_count, | |
| "total": len(models), | |
| "startupMs": startup_ms, | |
| "runtime": runtime, | |
| "workers": dispatcher.workers, | |
| # Our own capability, NOT the negotiated minimum: the Node side runs | |
| # the same min() against what it declared, so both ends derive the | |
| # identical effective version (D350). | |
| "protocolVersion": PROTOCOL_VERSION, | |
| }) | |
| # ── Window accumulator (per-model) ────────────────────────────── | |
| # | |
| # In `window` mode, every MSG_INFER_RAW arrival pushes (req_id, img) | |
| # into a per-model pending queue and arms a `window_ms` timer. When | |
| # the timer fires, OR the queue reaches `max_batch_size`, we flush | |
| # via `dispatcher.run_list(slot, imgs)` — single batched predict. | |
| # Concurrent inferRaw calls from N cameras coalesce into one CoreML | |
| # predict call, amortising ANE setup. | |
| class WindowAccumulator: | |
| def __init__(self, model_idx: int) -> None: | |
| self.model_idx = model_idx | |
| self.pending: list[tuple[int, Image.Image]] = [] | |
| self.flush_task: Optional[asyncio.Task[None]] = None | |
| async def submit(self, req_id: int, img: Image.Image) -> None: | |
| self.pending.append((req_id, img)) | |
| if len(self.pending) >= max_batch_size: | |
| # Reached cap — cancel any pending timer (we are flushing | |
| # immediately from the caller's coroutine, not the timer) | |
| # and flush now. | |
| if self.flush_task is not None: | |
| self.flush_task.cancel() | |
| self.flush_task = None | |
| await self._flush() | |
| return | |
| if self.flush_task is None: | |
| self.flush_task = asyncio.create_task(self._delayed_flush()) | |
| async def _delayed_flush(self) -> None: | |
| try: | |
| await asyncio.sleep(window_ms / 1000.0) | |
| except asyncio.CancelledError: | |
| return | |
| # Mark the timer task as complete BEFORE the dispatch await so | |
| # `_flush` does not try to cancel the running task (which would | |
| # raise CancelledError at `await dispatcher.run_list` and drop | |
| # responses on the floor — the symptom the TS side observed as | |
| # `runPipeline.mutate` hanging forever). | |
| self.flush_task = None | |
| await self._flush() | |
| async def _flush(self) -> None: | |
| batch = self.pending | |
| self.pending = [] | |
| if not batch: | |
| return | |
| slot = models[self.model_idx] | |
| try: | |
| imgs = [b[1] for b in batch] | |
| results = await dispatcher.run_list(slot, imgs) | |
| except Exception as exc: | |
| sys.stderr.write(f"Window flush error (model {self.model_idx}): {exc}\n") | |
| sys.stderr.flush() | |
| err_payload = { | |
| "error": str(exc), | |
| "kind": "detections", | |
| "detections": [], | |
| "inferenceMs": 0, | |
| } | |
| for rid, _ in batch: | |
| await writer.send(rid, err_payload) | |
| return | |
| for (rid, _), result in zip(batch, results): | |
| await writer.send(rid, result) | |
| accumulators: dict[int, WindowAccumulator] = {} | |
| # ── Bench frame cache ────────────────────────────────────────── | |
| # Stores raw PIL Images keyed by uint32 frame_id. The Node side | |
| # sends MSG_CACHE_FRAME once with the full 1.2MB payload; all | |
| # subsequent MSG_INFER_CACHED calls send only a 5-byte header | |
| # (model_idx + frame_id). Eliminates the 35ms/call pipe transfer | |
| # that dominates bench throughput. | |
| frame_cache: dict[int, Image.Image] = {} | |
| # Preprocessed-input cache: keyed by (frame_id, model_idx). First | |
| # inferCached call preprocesses and caches; subsequent calls skip | |
| # preprocess entirely and go straight to predict + postprocess. | |
| # Saves ~10ms/call (the PIL resize + numpy + tensor pack cost). | |
| _preprocess_cache: dict[tuple[int, int], tuple[dict, float, tuple[int, int], int, int]] = {} | |
| def inline_command(cmd: dict) -> dict: | |
| # Needs the loop's own state (caches, backpressure, counters), which | |
| # the model-command path never sees. Answered at once — this is the | |
| # host's liveness probe, and it must not wait for a compile. | |
| if cmd.get("cmd") == "mem_stats": | |
| return _mem_stats_payload( | |
| models, | |
| counters, | |
| backpressure, | |
| len(frame_cache), | |
| len(_preprocess_cache), | |
| len(_bench_preprocess_cache), | |
| round(time.monotonic() - pool_started_monotonic), | |
| model_load=model_loader.snapshot(), | |
| ) | |
| fid = cmd.get("frameId", -1) | |
| frame_cache.pop(fid, None) | |
| # Purge preprocessed tensor caches for this frame | |
| for k in [k for k in _preprocess_cache if k[0] == fid]: | |
| del _preprocess_cache[k] | |
| # Also purge the bench preprocess cache | |
| for k in [k for k in _bench_preprocess_cache if k[0] == fid]: | |
| del _bench_preprocess_cache[k] | |
| return {"cmd": "uncache_frame", "status": "ok", "frameId": fid} | |
| command_router = CommandRouter(models, model_loader, writer.send, inline_command) | |
| def get_accumulator(model_idx: int) -> WindowAccumulator: | |
| acc = accumulators.get(model_idx) | |
| if acc is None: | |
| acc = WindowAccumulator(model_idx) | |
| accumulators[model_idx] = acc | |
| return acc | |
| # ── Main loop ─────────────────────────────────────────────────── | |
| async def handle_inference( | |
| req_id: int, img: Image.Image, model_idx: int, deadline_ms: int = 0, | |
| ) -> None: | |
| try: | |
| # The deadline gate lives in _execute_inference, immediately | |
| # before predict — it runs here on admission AND again when a | |
| # queued frame is promoted below, which is where a request goes | |
| # stale (D350). | |
| await _execute_inference( | |
| req_id, img, model_idx, deadline_ms, | |
| models=models, | |
| dispatcher=dispatcher, | |
| writer=writer, | |
| batch_mode=batch_mode, | |
| get_accumulator=get_accumulator, | |
| counters=counters, | |
| ) | |
| except Exception as exc: | |
| sys.stderr.write(f"Inference error (model {model_idx}, req {req_id}): {exc}\n") | |
| sys.stderr.flush() | |
| await writer.send(req_id, { | |
| "error": str(exc), | |
| "kind": "detections", | |
| "detections": [], | |
| "inferenceMs": 0, | |
| }) | |
| finally: | |
| # Release this model's running slot; promote the next queued | |
| # frame (if any) into the freed slot immediately. | |
| for next_req_id, next_img, next_deadline_ms in backpressure.complete(model_idx): | |
| asyncio.create_task( | |
| handle_inference(next_req_id, next_img, model_idx, next_deadline_ms) | |
| ) | |
| def _send_dropped(req_id: int) -> None: | |
| # Fast shed response — returned in microseconds instead of queuing the | |
| # frame for seconds. The TS side (shared-inference-pool.ts) recognises | |
| # `dropped: true` and counts it, so overload is measurable, not silent. | |
| asyncio.create_task(writer.send(req_id, { | |
| "kind": "detections", | |
| "detections": [], | |
| "inferenceMs": 0, | |
| "dropped": True, | |
| })) | |
| def _dispatch_inference( | |
| req_id: int, img: Image.Image, model_idx: int, deadline_ms: int = 0, | |
| ) -> None: | |
| # Gate every single-frame inference through the per-model in-flight | |
| # bound. Slot accounting is synchronous (admit here, complete in the | |
| # handle_inference `finally`), so a slot can never leak past a task's | |
| # lifetime. Under overload the OLDEST queued frame is shed and answered | |
| # immediately so the caller never hangs on a frame we chose not to run. | |
| # The wire deadline rides the queued tuple so the gate in | |
| # _execute_inference sees it AFTER the wait, not just at arrival. | |
| to_run, to_drop = backpressure.admit(model_idx, (req_id, img, deadline_ms)) | |
| for dropped_req_id, _dropped_img, _dropped_deadline in to_drop: | |
| _send_dropped(dropped_req_id) | |
| for run_req_id, run_img, run_deadline_ms in to_run: | |
| asyncio.create_task(handle_inference(run_req_id, run_img, model_idx, run_deadline_ms)) | |
| async def handle_batch( | |
| req_id: int, model_idx: int, items: list[Image.Image], deadline_ms: int = 0, | |
| ) -> None: | |
| # Deadline gate first — the whole batch is one abandoned caller, so | |
| # every item is answered as a well-formed dropped record and none of | |
| # the work runs (D350). | |
| now_ms = time.time() * 1000.0 | |
| if _deadline_expired(deadline_ms, now_ms): | |
| counters.shed_deadline_expired += len(items) | |
| _log_deadline_shed(counters.shed_deadline_expired, model_idx, now_ms - deadline_ms) | |
| await writer.send(req_id, { | |
| "results": [_dropped_deadline_payload() for _ in items], | |
| }) | |
| return | |
| if model_idx >= len(models) or not models[model_idx].loaded: | |
| await writer.send(req_id, { | |
| "error": f"Model {model_idx} not loaded", | |
| "results": [], | |
| }) | |
| return | |
| try: | |
| slot = models[model_idx] | |
| if batch_mode in ("list", "window"): | |
| # `list` mode: send the batch through dispatcher.run_list | |
| # — single predict call, internal CoreML iteration. | |
| # `window` mode: same dispatch path (no point queueing into | |
| # the accumulator when caller already packed N items). | |
| payload = await dispatcher.run_list(slot, items) | |
| else: | |
| tasks = [dispatcher.run(slot, img) for img in items] | |
| results = await asyncio.gather(*tasks, return_exceptions=True) | |
| payload = [] | |
| for r in results: | |
| if isinstance(r, Exception): | |
| payload.append({ | |
| "error": str(r), | |
| "kind": "detections", | |
| "detections": [], | |
| "inferenceMs": 0, | |
| }) | |
| else: | |
| payload.append(r) | |
| await writer.send(req_id, {"results": payload}) | |
| except Exception as exc: | |
| sys.stderr.write(f"Batch error (model {model_idx}, req {req_id}): {exc}\n") | |
| sys.stderr.flush() | |
| await writer.send(req_id, {"error": str(exc), "results": []}) | |
| while True: | |
| try: | |
| header = await _read_exact(reader, PREFIX_LEN) | |
| except asyncio.IncompleteReadError: | |
| return # stdin closed | |
| total_len, req_id = struct.unpack("<II", header[0:8]) | |
| msg_type = header[8] | |
| remaining = total_len - 5 | |
| try: | |
| payload = await _read_exact(reader, remaining) if remaining > 0 else b"" | |
| except asyncio.IncompleteReadError: | |
| return | |
| # v2 sheddable payloads carry [8B absolute deadline] first; v1 (and | |
| # every command) passes through untouched with deadline 0 (D350). | |
| try: | |
| deadline_ms, payload = _split_deadline(msg_type, payload, wire_version) | |
| except ValueError as exc: | |
| await writer.send(req_id, {"error": str(exc)}) | |
| continue | |
| if msg_type == MSG_COMMAND: | |
| counters.commands += 1 | |
| try: | |
| cmd = json.loads(payload) | |
| except Exception as exc: | |
| await writer.send(req_id, {"cmd": "unknown", "status": "error", "error": str(exc)}) | |
| continue | |
| # NEVER awaits a model load: `load`/`replace` are queued and | |
| # compiled on their own thread under a bound (D653), so the next | |
| # frame is read while a GPU compile runs. | |
| await command_router.submit(req_id, cmd) | |
| elif msg_type == MSG_INFER_JPEG: | |
| counters.infer_jpeg += 1 | |
| if len(payload) < 1: | |
| await writer.send(req_id, {"error": "empty infer_jpeg payload"}) | |
| continue | |
| model_idx = payload[0] | |
| jpeg = payload[1:] | |
| try: | |
| img = decode_jpeg(jpeg) | |
| except Exception as exc: | |
| await writer.send(req_id, {"error": f"jpeg decode failed: {exc}"}) | |
| continue | |
| _dispatch_inference(req_id, img, model_idx, deadline_ms) | |
| elif msg_type == MSG_INFER_RAW: | |
| counters.infer_raw += 1 | |
| if len(payload) < 10: | |
| await writer.send(req_id, {"error": "truncated infer_raw header"}) | |
| continue | |
| model_idx = payload[0] | |
| width, height = struct.unpack("<II", payload[1:9]) | |
| fmt = payload[9] | |
| raw = payload[10:] | |
| try: | |
| img = wrap_raw(raw, width, height, fmt) | |
| except Exception as exc: | |
| await writer.send(req_id, {"error": f"raw wrap failed: {exc}"}) | |
| continue | |
| _dispatch_inference(req_id, img, model_idx, deadline_ms) | |
| elif msg_type == MSG_INFER_BATCH: | |
| # Header: [1B model_idx][1B count][4B frame_id] | |
| # frame_id (0 = untagged, the live/default case). A NONZERO frame_id | |
| # tags every wrapped item with `_bench_frame_id` so `_preprocess` | |
| # uses the bench cache — the whole batch is the SAME pinned frame in | |
| # the sustained-throughput benchmark, so decode+preprocess runs ONCE | |
| # (first miss) and every later inference is a pure-inference cache | |
| # hit. Historically the batch path shipped no frame_id and re-ran the | |
| # full preprocess on every inference (the ~78fps ceiling). | |
| counters.infer_batch += 1 | |
| if len(payload) < 6: | |
| await writer.send(req_id, {"error": "truncated infer_batch header"}) | |
| continue | |
| model_idx = payload[0] | |
| count = payload[1] | |
| batch_frame_id = struct.unpack("<I", payload[2:6])[0] | |
| offset = 6 | |
| items: list[Image.Image] = [] | |
| parse_err: Optional[str] = None | |
| for _ in range(count): | |
| # Per-item header: [4B w][4B h][1B fmt][4B size] | |
| if offset + 13 > len(payload): | |
| parse_err = "truncated batch item header" | |
| break | |
| width, height = struct.unpack("<II", payload[offset:offset + 8]) | |
| fmt = payload[offset + 8] | |
| size = struct.unpack("<I", payload[offset + 9:offset + 13])[0] | |
| offset += 13 | |
| if offset + size > len(payload): | |
| parse_err = "truncated batch item payload" | |
| break | |
| raw = payload[offset:offset + size] | |
| offset += size | |
| try: | |
| item_img = wrap_raw(raw, width, height, fmt) | |
| if batch_frame_id != 0: | |
| item_img._bench_frame_id = batch_frame_id | |
| items.append(item_img) | |
| except Exception as exc: | |
| parse_err = f"raw wrap failed: {exc}" | |
| break | |
| if parse_err is not None: | |
| await writer.send(req_id, {"error": parse_err, "results": []}) | |
| continue | |
| asyncio.create_task(handle_batch(req_id, model_idx, items, deadline_ms)) | |
| elif msg_type == MSG_CACHE_FRAME: | |
| counters.frames_cached += 1 | |
| # Wire: [4B frame_id][4B width][4B height][1B fmt][raw bytes] | |
| if len(payload) < 13: | |
| await writer.send(req_id, {"error": "truncated cache_frame header"}) | |
| continue | |
| fid = struct.unpack("<I", payload[0:4])[0] | |
| width, height = struct.unpack("<II", payload[4:12]) | |
| fmt = payload[12] | |
| raw = payload[13:] | |
| try: | |
| img = wrap_raw(raw, width, height, fmt) | |
| img._bench_frame_id = fid # Tag for preprocess cache keying | |
| frame_cache[fid] = img | |
| await writer.send(req_id, {"status": "cached", "frameId": fid, "width": width, "height": height}) | |
| except Exception as exc: | |
| await writer.send(req_id, {"error": f"cache_frame wrap failed: {exc}"}) | |
| elif msg_type == MSG_INFER_CACHED: | |
| # Wire: [1B model_idx][4B frame_id] | |
| # Routes through handle_inference → window accumulator for | |
| # batching. The preprocess runs each time (~4ms) but batch=4 | |
| # predict saves much more (8ms vs 4×9ms=36ms). | |
| counters.infer_cached += 1 | |
| if len(payload) < 5: | |
| await writer.send(req_id, {"error": "truncated infer_cached header"}) | |
| continue | |
| model_idx = payload[0] | |
| fid = struct.unpack("<I", payload[1:5])[0] | |
| img = frame_cache.get(fid) | |
| if img is None: | |
| await writer.send(req_id, {"error": f"frame {fid} not in cache"}) | |
| continue | |
| if model_idx >= len(models) or not models[model_idx].loaded: | |
| await writer.send(req_id, {"error": f"model {model_idx} not loaded"}) | |
| continue | |
| _dispatch_inference(req_id, img, model_idx, deadline_ms) | |
| else: | |
| await writer.send(req_id, {"error": f"unknown msg_type: {msg_type}"}) | |
| def main() -> None: | |
| try: | |
| asyncio.run(_run()) | |
| except KeyboardInterrupt: | |
| pass | |
| if __name__ == "__main__": | |
| main() | |