sevenhill/Immich
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: RKNN half RAM model usage + no weight duplication in multi-core (#23958)

Browse Source 
* perf(ml): share RKNN weights across NPU cores and halve model RAM

* perf(ml): default rknn threads to npu cores and cache shape setup

* perf(ml): serve rknn runs in arrival order

(cherry picked from commit e659fb62f2bb3f6b1db08ed643d3d054a539c415)
This commit is contained in:
Horatiu George Todoran
2026-10-28 20:40:53 +03:00
committed by Codex
parent 1f422ea4c4
commit f418967d55
9 changed files with 468 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
machine-learning/immich_ml/config.py
+1 -1
View File
@@ -70,7 +70,7 @@ class Settings(BaseSettings):
ann_fp16_turbo: bool = False
ann_tuning_level: int = 2
rknn: bool = True
rknn_threads: int = 1
rknn_threads: int | None = None
preload: PreloadModelData | None = None
max_batch_size: MaxBatchSize = MaxBatchSize()
model_organization: ModelOrganization = ModelOrganization.APP
machine-learning/immich_ml/models/constants.py
+1 -2
View File
@@ -100,8 +100,7 @@ SUPPORTED_PROVIDERS = [
"CPUExecutionProvider",
]
RKNN_SUPPORTED_SOCS = ["rk3566", "rk3568", "rk3576", "rk3588"]
RKNN_COREMASK_SUPPORTED_SOCS = ["rk3576", "rk3588"]
RKNN_NPU_CORES = {"rk3566": 1, "rk3568": 1, "rk3576": 2, "rk3588": 3}
WEBLATE_TO_FLORES200 = {
machine-learning/immich_ml/sessions/rknn/__init__.py
+6 -11
View File
@@ -9,9 +9,10 @@ import orjson
from numpy.typing import NDArray
from immich_ml.config import log, settings
from immich_ml.models.constants import RKNN_NPU_CORES
from immich_ml.schemas import ModelInput, ModelTensor, SessionNode, Shape
from .rknnpool import RknnNode, RknnPoolExecutor, is_available, native_outputs, soc_name
from .rknnpool import RknnNode, RknnPoolExecutor, is_available, soc_name
is_available = is_available and settings.rknn
model_prefix = Path("rknpu") / soc_name if is_available and soc_name is not None else None
@@ -21,12 +22,6 @@ def model_path(model_dir: Path, variant: str = "") -> Path:
return (model_dir / model_prefix if model_prefix else model_dir) / variant / "model.rknn"
def run_inference(rknn_lite: Any, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
rknn_lite.rknn_runtime.set_inputs(inputs, None, data_format)
rknn_lite.rknn_runtime.run(False)
return native_outputs(rknn_lite)
def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None:
if array.ndim != 4:
return None # token ids carry no layout
@@ -39,11 +34,11 @@ def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None:
class RknnSession:
def __init__(self, model_path: Path) -> None:
self.tpe = settings.rknn_threads
self.tpe = settings.rknn_threads or RKNN_NPU_CORES.get(soc_name or "", 1)
log.info(f"Loading RKNN model from {model_path} with {self.tpe} threads.")
self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe, func=run_inference)
log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads.")
self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe)
log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads sharing its weights.")
# the shapes the binary was compiled for, among which the runtime routes itself
batch = self.rknnpool.inputs[0].shape[0]
declared = orjson.loads(self.rknnpool.custom_string) if self.rknnpool.custom_string else {}
@@ -77,7 +72,7 @@ class RknnSession:
input_feed: ModelInput,
run_options: Any = None,
) -> list[NDArray[np.float32]]:
inputs = [array if array.flags.c_contiguous else array.copy() for array in input_feed.values()]
inputs = list(input_feed.values())
return self.rknnpool.run(inputs, input_layout(self.rknnpool.inputs[0].shape, inputs[0]))
machine-learning/immich_ml/sessions/rknn/rknnpool.py
+199 -94
View File
@@ -1,18 +1,15 @@
# This code is from leafqycc/rknn-multi-threaded
# Following Apache License 2.0
import ctypes
import logging
import threading
from concurrent.futures import ThreadPoolExecutor
from ctypes import byref, sizeof
from pathlib import Path
from typing import Any, Callable, NamedTuple
from queue import SimpleQueue
from typing import Any, NamedTuple
import numpy as np
from numpy.typing import NDArray
from immich_ml.config import log
from immich_ml.models.constants import RKNN_COREMASK_SUPPORTED_SOCS, RKNN_SUPPORTED_SOCS
from immich_ml.models.constants import RKNN_NPU_CORES
from immich_ml.schemas import ModelTensor
@@ -20,7 +17,7 @@ def get_soc(device_tree_path: Path | str) -> str | None:
try:
with Path(device_tree_path).open() as f:
device_compatible_str = f.read()
for soc in RKNN_SUPPORTED_SOCS:
for soc in RKNN_NPU_CORES:
if soc in device_compatible_str:
return soc
log.warning("Device is not supported for RKNN")
@@ -29,26 +26,34 @@ def get_soc(device_tree_path: Path | str) -> str | None:
return None
soc_name = None
is_available = False
try:
from rknnlite.api import RKNNLite, rknn_runtime
soc_name = get_soc("/proc/device-tree/compatible")
is_available = soc_name is not None
except ImportError:
log.debug("RKNN is not available")
class RknnNode(NamedTuple):
name: str
shape: tuple[int, ...]
_CUSTOM_STRING_QUERY = 7 # RKNN_QUERY_CUSTOM_STRING
# RKNN_QUERY_OUTPUT_ATTR, and RKNN_QUERY_CURRENT_OUTPUT_ATTR for the shape a dynamic binary last produced
_OUTPUT_QUERY, _CURRENT_OUTPUT_QUERY = 2, 15
_DTYPES = {0: np.float32, 1: np.float16} # rknn_tensor_type
# rknn_query_cmd
_IN_OUT_NUM, _INPUT_ATTR, _OUTPUT_ATTR, _CUSTOM_STRING = 0, 1, 2, 7
_INPUT_DYNAMIC_RANGE, _CURRENT_OUTPUT_ATTR = 13, 15
# rknn_tensor_format
_NCHW, _NHWC = 0, 1
# rknn_tensor_type, in enum order
_DTYPES: list[type[np.generic]] = [
np.float32,
np.float16,
np.int8,
np.uint8,
np.int16,
np.uint16,
np.int32,
np.uint32,
np.int64,
np.bool_,
]
_TYPES = {np.dtype(dtype): index for index, dtype in enumerate(_DTYPES)}
class _IoNum(ctypes.Structure):
_fields_ = [("n_input", ctypes.c_uint32), ("n_output", ctypes.c_uint32)]
class _CustomString(ctypes.Structure):
@@ -76,94 +81,194 @@ class _TensorAttr(ctypes.Structure): # rknn_tensor_attr
]
def native_outputs(rknn_lite: "RKNNLite") -> list[NDArray[Any]]:
"""The outputs in the precision the NPU computed them in, which RKNNLite would widen to fp32 on the CPU."""
runtime: Any = rknn_lite.rknn_runtime
count = runtime.get_in_out_num()[1]
buffers = (rknn_runtime.RKNNOutput * count)()
for index, buffer in enumerate(buffers):
buffer.index = index
if runtime.lib.rknn_outputs_get(runtime.context, count, buffers, None) != 0:
raise RuntimeError("RKNN inference failed!")
query = _CURRENT_OUTPUT_QUERY if runtime.is_dynamic_shape else _OUTPUT_QUERY
outputs = []
for buffer in buffers:
attr = _TensorAttr(index=buffer.index)
runtime.lib.rknn_query(runtime.context, query, ctypes.byref(attr), ctypes.sizeof(attr))
data = ctypes.cast(buffer.buf, ctypes.POINTER(ctypes.c_char * buffer.size)).contents
outputs.append(np.frombuffer(data, _DTYPES[attr.type]).reshape(attr.dims[: attr.n_dims]).copy())
runtime.lib.rknn_outputs_release(runtime.context, count, buffers)
return outputs
class _InputRange(ctypes.Structure):
_fields_ = [
("index", ctypes.c_uint32),
("shape_number", ctypes.c_uint32),
("fmt", ctypes.c_int),
("name", ctypes.c_char * 256),
("dyn_range", (ctypes.c_uint32 * 16) * 512),
("n_dims", ctypes.c_uint32),
]
def tensor_nodes(rknn_lite: "RKNNLite") -> tuple[list[RknnNode], list[RknnNode]]:
runtime = rknn_lite.rknn_runtime
n_inputs, n_outputs = runtime.get_in_out_num()
def node(index: int, is_output: bool) -> RknnNode:
attr = runtime.get_tensor_attr(index, is_output=is_output)
return RknnNode(attr.name.decode(), tuple(attr.dims[: attr.n_dims]))
return [node(i, False) for i in range(n_inputs)], [node(i, True) for i in range(n_outputs)]
class _Input(ctypes.Structure):
_fields_ = [
("index", ctypes.c_uint32),
("buf", ctypes.c_void_p),
("size", ctypes.c_uint32),
("pass_through", ctypes.c_uint8),
("type", ctypes.c_int),
("fmt", ctypes.c_int),
]
def custom_string(rknn_lite: "RKNNLite") -> str:
runtime: Any = rknn_lite.rknn_runtime
buffer = _CustomString()
# rknnlite declares argtypes on the shared handle, so pass plain values and let ctypes convert
runtime.lib.rknn_query(runtime.context, _CUSTOM_STRING_QUERY, ctypes.byref(buffer), ctypes.sizeof(buffer))
declared: str = buffer.string.decode()
return declared
class _Output(ctypes.Structure):
_fields_ = [
("want_float", ctypes.c_uint8),
("is_prealloc", ctypes.c_uint8),
("index", ctypes.c_uint32),
("buf", ctypes.c_void_p),
("size", ctypes.c_uint32),
]
def init_rknn(model_path: str) -> "RKNNLite":
if not is_available:
raise RuntimeError("rknn is not available!")
rknn_lite = RKNNLite()
rknn_lite.rknn_log.logger.setLevel(logging.ERROR)
ret = rknn_lite.load_rknn(model_path)
_Key = tuple[tuple[int, ...], ...]
_Spec = tuple[tuple[int, ...], np.dtype[Any]]
_context = ctypes.c_uint64 # rknn_context on 64-bit
_ctx_p = ctypes.POINTER(_context)
_u32 = ctypes.c_uint32
soc_name = None
is_available = False
try:
lib = ctypes.CDLL("librknnrt.so") # CDLL releases the GIL for each call, so contexts run in parallel
lib.rknn_init.argtypes = _ctx_p, ctypes.c_char_p, _u32, _u32, ctypes.c_void_p
lib.rknn_dup_context.argtypes = _ctx_p, _ctx_p
lib.rknn_destroy.argtypes = (_context,)
lib.rknn_query.argtypes = _context, ctypes.c_int, ctypes.c_void_p, _u32
lib.rknn_set_input_shapes.argtypes = _context, _u32, ctypes.POINTER(_TensorAttr)
lib.rknn_inputs_set.argtypes = _context, _u32, ctypes.POINTER(_Input)
lib.rknn_run.argtypes = _context, ctypes.c_void_p
lib.rknn_outputs_get.argtypes = _context, _u32, ctypes.POINTER(_Output), ctypes.c_void_p
lib.rknn_outputs_release.argtypes = _context, _u32, ctypes.POINTER(_Output)
soc_name = get_soc("/proc/device-tree/compatible")
is_available = soc_name is not None
except OSError:
log.debug("RKNN is not available")
def _check(ret: int, call: str) -> None:
if ret != 0:
raise RuntimeError("Failed to load RKNN model")
raise RuntimeError(f"{call} failed with RKNN error {ret}")
if soc_name in RKNN_COREMASK_SUPPORTED_SOCS:
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
else:
ret = rknn_lite.init_runtime() # Please do not set this parameter on other platforms.
if ret != 0:
raise RuntimeError("Failed to initialize RKNN runtime environment")
def _query(ctx: int, cmd: int, info: ctypes.Structure) -> None:
_check(lib.rknn_query(ctx, cmd, byref(info), sizeof(info)), "rknn_query")
return rknn_lite
def _attr(ctx: int, cmd: int, index: int) -> _TensorAttr:
attr = _TensorAttr(index=index)
_query(ctx, cmd, attr)
return attr
def _specs(attrs: list[_TensorAttr]) -> list[_Spec]:
return [(tuple(attr.dims[: attr.n_dims]), np.dtype(_DTYPES[attr.type])) for attr in attrs]
class RknnPoolExecutor:
def __init__(
self,
model_path: str,
tpes: int,
func: Callable[["RKNNLite", list[ModelTensor], str | None], list[NDArray[np.float32]]],
) -> None:
self.tpes = tpes
self.rknn_pool = [init_rknn(model_path) for _ in range(tpes)]
self.inputs, self.outputs = tensor_nodes(self.rknn_pool[0])
self.custom_string = custom_string(self.rknn_pool[0])
self.pool = ThreadPoolExecutor(max_workers=tpes)
self.func = func
self.num = 0
self.lock = threading.Lock()
"""One context per thread, all but the first duplicated from it so that they share its weights.
def run(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
with self.lock:
idx = self.num % self.tpes
self.num += 1
The model is loaded by path rather than from a buffer, so the runtime holds the only copy of it."""
fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs, data_format)
return fut.result()
def __init__(self, model_path: str, tpes: int) -> None:
self.contexts: list[int] = []
self.executor = ThreadPoolExecutor(tpes, thread_name_prefix="rknn")
master = _context()
_check(lib.rknn_init(byref(master), model_path.encode(), 0, 0, None), "rknn_init")
self.contexts.append(master.value)
for _ in range(tpes - 1):
child = _context()
_check(lib.rknn_dup_context(byref(master), byref(child)), "rknn_dup_context")
self.contexts.append(child.value)
self.idle: SimpleQueue[int] = SimpleQueue()
for ctx in self.contexts:
self.idle.put(ctx)
io = _IoNum()
_query(master.value, _IN_OUT_NUM, io)
self.input_attrs = [_attr(master.value, _INPUT_ATTR, i) for i in range(io.n_input)]
self.output_attrs = [_attr(master.value, _OUTPUT_ATTR, i) for i in range(io.n_output)]
self.inputs = [RknnNode(a.name.decode(), tuple(a.dims[: a.n_dims])) for a in self.input_attrs]
self.outputs = [RknnNode(a.name.decode(), tuple(a.dims[: a.n_dims])) for a in self.output_attrs]
custom = _CustomString()
lib.rknn_query(master.value, _CUSTOM_STRING, byref(custom), sizeof(custom))
self.custom_string = custom.string.decode()
ranges = [_InputRange(index=index) for index in range(io.n_input)]
for rng in ranges:
lib.rknn_query(master.value, _INPUT_DYNAMIC_RANGE, byref(rng), sizeof(rng))
self.formats = [rng.fmt for rng in ranges]
self.compiled = [{tuple(rng.dyn_range[i][: rng.n_dims]) for i in range(rng.shape_number)} for rng in ranges]
self.dynamic = any(self.compiled)
self.shape_attrs: dict[_Key, ctypes.Array[_TensorAttr]] = {}
self.output_specs: dict[_Key, list[_Spec]] = {(): _specs(self.output_attrs)}
self.current: dict[int, _Key] = {}
def run(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[Any]]:
return self.executor.submit(self._next, inputs, data_format).result()
def _next(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[Any]]:
ctx = self.idle.get()
try:
return self._run(ctx, inputs, data_format)
finally:
self.idle.put(ctx)
def _run(self, ctx: int, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[Any]]:
fed: list[NDArray[Any]] = [
np.ascontiguousarray(array.transpose(0, 2, 3, 1) if data_format == "nchw" and array.ndim == 4 else array)
for array in inputs
]
key = self._set_shapes(ctx, fed) if self.dynamic else ()
feed = (_Input * len(fed))(
*(
_Input(index=i, buf=array.ctypes.data, size=array.nbytes, type=_TYPES[array.dtype], fmt=_NHWC)
for i, array in enumerate(fed)
)
)
_check(lib.rknn_inputs_set(ctx, len(fed), feed), "rknn_inputs_set")
_check(lib.rknn_run(ctx, None), "rknn_run")
specs = self.output_specs.get(key)
if specs is None:
attrs = [_attr(ctx, _CURRENT_OUTPUT_ATTR, i) for i in range(len(self.output_attrs))]
specs = self.output_specs[key] = _specs(attrs)
outputs = [np.empty(dims, dtype) for dims, dtype in specs]
buffers = (_Output * len(outputs))(
*(
_Output(is_prealloc=1, index=i, buf=array.ctypes.data, size=array.nbytes)
for i, array in enumerate(outputs)
)
)
_check(lib.rknn_outputs_get(ctx, len(outputs), buffers, None), "rknn_outputs_get")
lib.rknn_outputs_release(ctx, len(outputs), buffers)
return outputs
def _set_shapes(self, ctx: int, inputs: list[NDArray[Any]]) -> _Key:
key = tuple(
self._compiled_shape(fmt, compiled, array.shape)
for fmt, compiled, array in zip(self.formats, self.compiled, inputs)
)
if self.current.get(ctx) == key:
return key
attrs = self.shape_attrs.get(key)
if attrs is None:
attrs = self.shape_attrs[key] = (_TensorAttr * len(key))()
for i, (source, compiled, fmt, shape) in enumerate(zip(self.input_attrs, self.compiled, self.formats, key)):
attrs[i] = _TensorAttr.from_buffer_copy(source)
if compiled:
attrs[i].n_dims = len(shape)
attrs[i].dims[: len(shape)] = shape
attrs[i].fmt = fmt
_check(lib.rknn_set_input_shapes(ctx, len(key), attrs), "rknn_set_input_shapes")
self.current[ctx] = key
return key
@staticmethod
def _compiled_shape(fmt: int, compiled: set[tuple[int, ...]], fed: tuple[int, ...]) -> tuple[int, ...]:
shape = (fed[0], fed[3], fed[1], fed[2]) if len(fed) == 4 and fmt == _NCHW else fed
if compiled and shape not in compiled:
raise ValueError(f"binary was compiled for {sorted(compiled)}, fed {list(fed)}")
return shape
def release(self) -> None:
self.pool.shutdown()
for rknn_lite in self.rknn_pool:
rknn_lite.release()
self.executor.shutdown()
for ctx in reversed(self.contexts): # duplicates before the context whose weights they share
lib.rknn_destroy(ctx)
self.contexts = []
def __del__(self) -> None:
self.release()
machine-learning/pyproject.toml
+1 -1
View File
@@ -57,7 +57,7 @@ cuda = [
]
openvino = ["onnxruntime-openvino>=1.24.1,<2"]
armnn = ["onnxruntime>=1.23.2,<2"]
rknn = ["onnxruntime>=1.23.2,<2", "rknn-toolkit-lite2>=2.3.0,<3"]
rknn = ["onnxruntime>=1.23.2,<2"]
rocm = ["onnxruntime-migraphx>=1.23.2,<2"]
[tool.uv]
machine-learning/test_main.py
+252 -4
View File
@@ -1,3 +1,4 @@
import ctypes
import json
import os
import pickle
@@ -63,8 +64,9 @@ from immich_ml.schemas import (
from immich_ml.sessions.ann import AnnSession
from immich_ml.sessions.ort import Device, GraphSpec, OrtSession, flush_denormals, fresh, prepared
from immich_ml.sessions.policy import ShapePolicy, batches, runs
from immich_ml.sessions.rknn import RknnSession, run_inference
from immich_ml.sessions.rknn import RknnSession
from immich_ml.sessions.rknn import model_path as rknn_model_path
from immich_ml.sessions.rknn.rknnpool import RknnPoolExecutor
class TestBase:
@@ -960,9 +962,11 @@ class TestRknnSession:
mocker.patch("immich_ml.sessions.rknn.is_available", True)
RknnSession(model_path)
rknn_session.assert_called_once_with(model_path=model_path.as_posix(), tpes=tpe, func=run_inference)
rknn_session.assert_called_once_with(model_path=model_path.as_posix(), tpes=tpe)
info.assert_has_calls([mock.call(f"Loaded RKNN model from {model_path} with {tpe} threads.")])
info.assert_has_calls(
[mock.call(f"Loaded RKNN model from {model_path} with {tpe} threads sharing its weights.")]
)
def test_run_rknn(self, rknn_session: mock.Mock, mocker: MockerFixture) -> None:
rknn_session.return_value.load.return_value = 123
@@ -974,7 +978,20 @@ class TestRknnSession:
session.run(None, input_feed)
rknn_session.return_value.run.assert_called_once_with([input1, input2], "nchw")
assert all(fed.flags.c_contiguous for fed in rknn_session.return_value.run.call_args.args[0])
@pytest.mark.parametrize(
("soc", "threads", "expected"),
[("rk3566", None, 1), ("rk3576", None, 2), ("rk3588", None, 3), ("rk3588", 2, 2)],
)
def test_uses_a_context_per_npu_core_unless_set(
self, soc: str, threads: int | None, expected: int, rknn_session: mock.Mock, mocker: MockerFixture
) -> None:
mocker.patch("immich_ml.sessions.rknn.soc_name", soc)
mocker.patch.object(settings, "rknn_threads", threads)
RknnSession(Path("ViT-B-32__openai"))
assert rknn_session.call_args.kwargs["tpes"] == expected
def test_run_rknn_rejects_a_shape_the_binary_was_not_compiled_for(
self, rknn_session: mock.Mock, mocker: MockerFixture
@@ -1011,6 +1028,237 @@ class TestRknnSession:
assert RknnSession(Path("buffalo_l")).shapes == (Shape(batch=1),)
class FakeRknn:
"""librknnrt for one input and one output, writing through the pointers as the runtime does."""
def __init__(self, ranges: list[tuple[int, ...]] | None = None) -> None:
self.ranges = ranges
self.contexts = 0
self.fed: list[tuple[int, int, int]] = [] # (context, type, fmt) per run
self.data: list[bytes] = []
self.buffers: list[int] = []
self.output_queries = 0
self.set_shapes: list[tuple[int, tuple[int, ...]]] = []
self.rknn_destroy = mock.Mock(return_value=0)
self.rknn_outputs_release = mock.Mock(return_value=0)
self.rknn_run = mock.Mock(return_value=0)
def rknn_init(self, ctx: Any, model: bytes, size: int, flag: int, extend: None) -> int:
self.contexts += 1
ctx._obj.value = self.contexts
return 0
def rknn_dup_context(self, ctx_in: Any, ctx_out: Any) -> int:
self.contexts += 1
ctx_out._obj.value = self.contexts
return 0
def rknn_query(self, ctx: int, cmd: int, info: Any, size: int) -> int:
obj = info._obj
if cmd == 0:
obj.n_input, obj.n_output = 1, 1
elif cmd == 1:
self._fill(obj, b"x", self.ranges[0] if self.ranges else (1, 48, 320, 3), 0)
obj.fmt = 1
elif cmd == 2:
self._fill(obj, b"logits", (1, 40, 97), 1)
elif cmd == 7:
obj.string = b'{"dims":[{}]}'
elif cmd == 13:
if self.ranges is None:
return -1
obj.shape_number, obj.n_dims, obj.fmt = len(self.ranges), 4, 1
for i, shape in enumerate(self.ranges):
obj.dyn_range[i][:4] = shape
elif cmd == 15:
self.output_queries += 1
width = dict(self.set_shapes)[ctx][2]
self._fill(obj, b"logits", (1, width // 8, 97), 1)
return 0
@staticmethod
def _fill(attr: Any, name: bytes, dims: tuple[int, ...], dtype: int) -> None:
attr.name, attr.n_dims, attr.type = name, len(dims), dtype
attr.dims[: len(dims)] = dims
def rknn_set_input_shapes(self, ctx: int, n: int, attrs: Any) -> int:
self.set_shapes.append((ctx, tuple(attrs[0].dims[: attrs[0].n_dims])))
return 0
def rknn_inputs_set(self, ctx: int, n: int, inputs: Any) -> int:
self.fed.append((ctx, inputs[0].type, inputs[0].fmt))
self.data.append(ctypes.string_at(inputs[0].buf, inputs[0].size))
self.buffers.append(inputs[0].buf)
return 0
def rknn_outputs_get(self, ctx: int, n: int, outputs: Any, extend: None) -> int:
assert outputs[0].is_prealloc and not outputs[0].want_float
np.frombuffer((ctypes.c_char * outputs[0].size).from_address(outputs[0].buf), np.uint8)[:] = 0
return 0
class TestRknnPool:
@pytest.fixture
def lib(self, mocker: MockerFixture) -> FakeRknn:
fake = FakeRknn()
mocker.patch("immich_ml.sessions.rknn.rknnpool.lib", fake, create=True)
return fake
@pytest.fixture
def dynamic_lib(self, mocker: MockerFixture) -> FakeRknn:
fake = FakeRknn(ranges=[(1, 48, 320, 3), (1, 48, 640, 3)])
mocker.patch("immich_ml.sessions.rknn.rknnpool.lib", fake, create=True)
return fake
def test_duplicates_one_context_per_thread(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 3)
assert pool.contexts == [1, 2, 3]
assert pool.inputs[0] == ("x", (1, 48, 320, 3))
assert pool.outputs[0] == ("logits", (1, 40, 97))
assert pool.custom_string == '{"dims":[{}]}'
assert not pool.dynamic
pool.release()
assert lib.rknn_destroy.call_args_list == [mock.call(3), mock.call(2), mock.call(1)]
def test_feeds_the_dtype_given_as_nhwc(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
image = np.arange(48 * 320 * 3, dtype=np.uint8).reshape(1, 48, 320, 3)
pool.run([image], "nhwc")
pool.run([image.transpose(0, 3, 1, 2).astype(np.float32)], "nchw")
assert lib.fed == [(1, 3, 1), (1, 0, 1)] # (context, rknn_tensor_type, rknn_tensor_format)
assert lib.data[0] == image.tobytes()
assert lib.data[1] == image.astype(np.float32).tobytes() # transposed back to channels last
def test_feeds_a_channels_first_view_of_channels_last_data_without_a_copy(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
image = np.zeros((1, 48, 320, 3), dtype=np.float32)
pool.run([image.transpose(0, 3, 1, 2)], "nchw")
assert lib.buffers == [image.ctypes.data]
def test_makes_every_input_contiguous(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
image = np.arange(48 * 640 * 3, dtype=np.uint8).reshape(1, 48, 640, 3)[:, :, ::2]
pool.run([image], "nhwc")
assert lib.data == [np.ascontiguousarray(image).tobytes()]
def test_returns_outputs_in_the_precision_computed(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
[output] = pool.run([np.zeros((1, 48, 320, 3), dtype=np.uint8)], "nhwc")
assert output.dtype == np.float16
assert output.shape == (1, 40, 97)
lib.rknn_outputs_release.assert_called_once()
def test_a_failed_init_destroys_nothing(self, lib: FakeRknn, mocker: MockerFixture) -> None:
mocker.patch.object(lib, "rknn_init", return_value=-6)
with pytest.raises(RuntimeError, match="rknn_init failed"):
RknnPoolExecutor("model.rknn", 3)
lib.rknn_destroy.assert_not_called()
def test_returns_the_context_when_a_run_fails(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 2)
lib.rknn_run.return_value = -1
with pytest.raises(RuntimeError, match="rknn_run failed"):
pool.run([np.zeros((1, 48, 320, 3), dtype=np.uint8)], "nhwc")
assert pool.idle.qsize() == 2
def test_runs_on_every_context_at_once(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 3)
barrier = threading.Barrier(3, timeout=5) # breaks unless all three are inside rknn_run together
lib.rknn_run.side_effect = lambda ctx, extend: barrier.wait() * 0
with ThreadPoolExecutor(3) as executor:
futures = [executor.submit(pool.run, [np.zeros((1, 48, 320, 3), dtype=np.uint8)], "nhwc") for _ in range(3)]
[future.result() for future in futures]
assert sorted(ctx for ctx, _, _ in lib.fed) == [1, 2, 3]
def test_a_waiting_run_is_not_overtaken_by_the_caller_that_just_ran(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
def run(ctx: int, extend: None) -> int:
time.sleep(0.01)
return 0
lib.rknn_run.side_effect = run
busy, waiting = np.zeros((1, 48, 320, 3), np.uint8), np.ones((1, 48, 320, 3), np.uint8)
caller = threading.Thread(target=lambda: [pool.run([busy], "nhwc") for _ in range(20)])
caller.start()
time.sleep(0.005)
pool.run([waiting], "nhwc")
caller.join(5)
assert [data[0] for data in lib.data].index(1) <= 2
def test_does_not_set_shapes_for_a_static_binary(self, lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
pool.run([np.zeros((1, 48, 320, 3), dtype=np.uint8)], "nhwc")
assert lib.set_shapes == []
def test_routes_a_dynamic_binary_to_the_shape_fed(self, dynamic_lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
[wide] = pool.run([np.zeros((1, 3, 48, 640), dtype=np.float32)], "nchw")
[narrow] = pool.run([np.zeros((1, 48, 320, 3), dtype=np.uint8)], "nhwc")
assert pool.dynamic
assert dynamic_lib.set_shapes == [(1, (1, 48, 640, 3)), (1, (1, 48, 320, 3))]
assert wide.shape == (1, 80, 97)
assert narrow.shape == (1, 40, 97)
def test_sets_the_shape_on_the_context_that_runs(self, dynamic_lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 2)
image = np.zeros((1, 48, 640, 3), dtype=np.uint8)
pool.run([image], "nhwc")
pool.run([image], "nhwc")
assert dynamic_lib.set_shapes == [(1, (1, 48, 640, 3)), (2, (1, 48, 640, 3))]
def test_sets_a_shape_again_only_when_it_changes(self, dynamic_lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
wide, narrow = np.zeros((1, 48, 640, 3), dtype=np.uint8), np.zeros((1, 48, 320, 3), dtype=np.uint8)
for image in (wide, wide, narrow, narrow, wide):
pool.run([image], "nhwc")
assert dynamic_lib.set_shapes == [(1, (1, 48, 640, 3)), (1, (1, 48, 320, 3)), (1, (1, 48, 640, 3))]
def test_queries_output_shapes_once_per_input_shape(self, dynamic_lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 2)
image = np.zeros((1, 48, 640, 3), dtype=np.uint8)
outputs = [pool.run([image], "nhwc")[0] for _ in range(4)]
assert dynamic_lib.output_queries == 1
assert all(output.shape == (1, 80, 97) for output in outputs)
def test_rejects_a_shape_the_binary_was_not_compiled_for(self, dynamic_lib: FakeRknn) -> None:
pool = RknnPoolExecutor("model.rknn", 1)
with pytest.raises(ValueError, match="compiled for"):
pool.run([np.zeros((1, 48, 480, 3), dtype=np.uint8)], "nhwc")
assert pool.idle.qsize() == 1
class TestCLIP:
embedding = np.random.rand(512).astype(np.float32)
cache_dir = Path("test_cache")
machine-learning/uv.lock
Generated
-53
View File
@@ -868,7 +868,6 @@ openvino = [
]
rknn = [
{ name = "onnxruntime" },
{ name = "rknn-toolkit-lite2" },
]
rocm = [
{ name = "onnxruntime-migraphx" },
@@ -943,7 +942,6 @@ requires-dist = [
{ name = "python-multipart", specifier = ">=0.0.6,<1.0" },
{ name = "rapidocr", specifier = ">=3.1.0" },
{ name = "rich", specifier = ">=13.4.2" },
{ name = "rknn-toolkit-lite2", marker = "extra == 'rknn'", specifier = ">=2.3.0,<3" },
{ name = "tokenizers", specifier = ">=0.15.0,<1.0" },
{ name = "uvicorn", extras = ["standard"], specifier = ">=0.22.0,<1.0" },
]
@@ -2303,57 +2301,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" },
]
[[package]]
name = "rknn-toolkit-lite2"
version = "2.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "numpy" },
{ name = "psutil" },
{ name = "ruamel-yaml" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/ff/db/64c756f3f06b219e92ff4f0fd4e000870ee49f214d505ff01c8b0275e26d/rknn_toolkit_lite2-2.3.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1e4ec691fed900c0e6fde5e7d8eeba17f806aa45092b63b361ee775e2c1b50e", size = 527458, upload-time = "2025-04-09T09:39:58.881Z" },
]
[[package]]
name = "ruamel-yaml"
version = "0.18.10"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "ruamel-yaml-clib", marker = "python_full_version < '3.13' and platform_python_implementation == 'CPython'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ea/46/f44d8be06b85bc7c4d8c95d658be2b68f27711f279bf9dd0612a5e4794f5/ruamel.yaml-0.18.10.tar.gz", hash = "sha256:20c86ab29ac2153f80a428e1254a8adf686d3383df04490514ca3b79a362db58", size = 143447, upload-time = "2025-01-06T14:08:51.334Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c2/36/dfc1ebc0081e6d39924a2cc53654497f967a084a436bb64402dfce4254d9/ruamel.yaml-0.18.10-py3-none-any.whl", hash = "sha256:30f22513ab2301b3d2b577adc121c6471f28734d3d9728581245f1e76468b4f1", size = 117729, upload-time = "2025-01-06T14:08:47.471Z" },
]
[[package]]
name = "ruamel-yaml-clib"
version = "0.2.12"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/20/84/80203abff8ea4993a87d823a5f632e4d92831ef75d404c9fc78d0176d2b5/ruamel.yaml.clib-0.2.12.tar.gz", hash = "sha256:6c8fbb13ec503f99a91901ab46e0b07ae7941cd527393187039aec586fdfd36f", size = 225315, upload-time = "2024-10-20T10:10:56.22Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/48/41/e7a405afbdc26af961678474a55373e1b323605a4f5e2ddd4a80ea80f628/ruamel.yaml.clib-0.2.12-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:20b0f8dc160ba83b6dcc0e256846e1a02d044e13f7ea74a3d1d56ede4e48c632", size = 133433, upload-time = "2024-10-20T10:12:55.657Z" },
{ url = "https://files.pythonhosted.org/packages/ec/b0/b850385604334c2ce90e3ee1013bd911aedf058a934905863a6ea95e9eb4/ruamel.yaml.clib-0.2.12-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:943f32bc9dedb3abff9879edc134901df92cfce2c3d5c9348f172f62eb2d771d", size = 647362, upload-time = "2024-10-20T10:12:57.155Z" },
{ url = "https://files.pythonhosted.org/packages/44/d0/3f68a86e006448fb6c005aee66565b9eb89014a70c491d70c08de597f8e4/ruamel.yaml.clib-0.2.12-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95c3829bb364fdb8e0332c9931ecf57d9be3519241323c5274bd82f709cebc0c", size = 754118, upload-time = "2024-10-20T10:12:58.501Z" },
{ url = "https://files.pythonhosted.org/packages/52/a9/d39f3c5ada0a3bb2870d7db41901125dbe2434fa4f12ca8c5b83a42d7c53/ruamel.yaml.clib-0.2.12-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:749c16fcc4a2b09f28843cda5a193e0283e47454b63ec4b81eaa2242f50e4ccd", size = 706497, upload-time = "2024-10-20T10:13:00.211Z" },
{ url = "https://files.pythonhosted.org/packages/b0/fa/097e38135dadd9ac25aecf2a54be17ddf6e4c23e43d538492a90ab3d71c6/ruamel.yaml.clib-0.2.12-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:bf165fef1f223beae7333275156ab2022cffe255dcc51c27f066b4370da81e31", size = 698042, upload-time = "2024-10-21T11:26:46.038Z" },
{ url = "https://files.pythonhosted.org/packages/ec/d5/a659ca6f503b9379b930f13bc6b130c9f176469b73b9834296822a83a132/ruamel.yaml.clib-0.2.12-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:32621c177bbf782ca5a18ba4d7af0f1082a3f6e517ac2a18b3974d4edf349680", size = 745831, upload-time = "2024-10-21T11:26:47.487Z" },
{ url = "https://files.pythonhosted.org/packages/db/5d/36619b61ffa2429eeaefaab4f3374666adf36ad8ac6330d855848d7d36fd/ruamel.yaml.clib-0.2.12-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b82a7c94a498853aa0b272fd5bc67f29008da798d4f93a2f9f289feb8426a58d", size = 715692, upload-time = "2024-12-11T19:58:17.252Z" },
{ url = "https://files.pythonhosted.org/packages/b1/82/85cb92f15a4231c89b95dfe08b09eb6adca929ef7df7e17ab59902b6f589/ruamel.yaml.clib-0.2.12-cp312-cp312-win32.whl", hash = "sha256:e8c4ebfcfd57177b572e2040777b8abc537cdef58a2120e830124946aa9b42c5", size = 98777, upload-time = "2024-10-20T10:13:01.395Z" },
{ url = "https://files.pythonhosted.org/packages/d7/8f/c3654f6f1ddb75daf3922c3d8fc6005b1ab56671ad56ffb874d908bfa668/ruamel.yaml.clib-0.2.12-cp312-cp312-win_amd64.whl", hash = "sha256:0467c5965282c62203273b838ae77c0d29d7638c8a4e3a1c8bdd3602c10904e4", size = 115523, upload-time = "2024-10-20T10:13:02.768Z" },
{ url = "https://files.pythonhosted.org/packages/29/00/4864119668d71a5fa45678f380b5923ff410701565821925c69780356ffa/ruamel.yaml.clib-0.2.12-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:4c8c5d82f50bb53986a5e02d1b3092b03622c02c2eb78e29bec33fd9593bae1a", size = 132011, upload-time = "2024-10-20T10:13:04.377Z" },
{ url = "https://files.pythonhosted.org/packages/7f/5e/212f473a93ae78c669ffa0cb051e3fee1139cb2d385d2ae1653d64281507/ruamel.yaml.clib-0.2.12-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:e7e3736715fbf53e9be2a79eb4db68e4ed857017344d697e8b9749444ae57475", size = 642488, upload-time = "2024-10-20T10:13:05.906Z" },
{ url = "https://files.pythonhosted.org/packages/1f/8f/ecfbe2123ade605c49ef769788f79c38ddb1c8fa81e01f4dbf5cf1a44b16/ruamel.yaml.clib-0.2.12-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b7e75b4965e1d4690e93021adfcecccbca7d61c7bddd8e22406ef2ff20d74ef", size = 745066, upload-time = "2024-10-20T10:13:07.26Z" },
{ url = "https://files.pythonhosted.org/packages/e2/a9/28f60726d29dfc01b8decdb385de4ced2ced9faeb37a847bd5cf26836815/ruamel.yaml.clib-0.2.12-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:96777d473c05ee3e5e3c3e999f5d23c6f4ec5b0c38c098b3a5229085f74236c6", size = 701785, upload-time = "2024-10-20T10:13:08.504Z" },
{ url = "https://files.pythonhosted.org/packages/84/7e/8e7ec45920daa7f76046578e4f677a3215fe8f18ee30a9cb7627a19d9b4c/ruamel.yaml.clib-0.2.12-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:3bc2a80e6420ca8b7d3590791e2dfc709c88ab9152c00eeb511c9875ce5778bf", size = 693017, upload-time = "2024-10-21T11:26:48.866Z" },
{ url = "https://files.pythonhosted.org/packages/c5/b3/d650eaade4ca225f02a648321e1ab835b9d361c60d51150bac49063b83fa/ruamel.yaml.clib-0.2.12-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:e188d2699864c11c36cdfdada94d781fd5d6b0071cd9c427bceb08ad3d7c70e1", size = 741270, upload-time = "2024-10-21T11:26:50.213Z" },
{ url = "https://files.pythonhosted.org/packages/87/b8/01c29b924dcbbed75cc45b30c30d565d763b9c4d540545a0eeecffb8f09c/ruamel.yaml.clib-0.2.12-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4f6f3eac23941b32afccc23081e1f50612bdbe4e982012ef4f5797986828cd01", size = 709059, upload-time = "2024-12-11T19:58:18.846Z" },
{ url = "https://files.pythonhosted.org/packages/30/8c/ed73f047a73638257aa9377ad356bea4d96125b305c34a28766f4445cc0f/ruamel.yaml.clib-0.2.12-cp313-cp313-win32.whl", hash = "sha256:6442cb36270b3afb1b4951f060eccca1ce49f3d087ca1ca4563a6eb479cb3de6", size = 98583, upload-time = "2024-10-20T10:13:09.658Z" },
{ url = "https://files.pythonhosted.org/packages/b0/85/e8e751d8791564dd333d5d9a4eab0a7a115f7e349595417fd50ecae3395c/ruamel.yaml.clib-0.2.12-cp313-cp313-win_amd64.whl", hash = "sha256:e5b8daf27af0b90da7bb903a876477a9e6d7270be6146906b276605997c7e9a3", size = 115190, upload-time = "2024-10-20T10:13:10.66Z" },
]
[[package]]
name = "ruff"
version = "0.17.1"