diff --git a/backend/fluksio/flow/nodes/wren.py b/backend/fluksio/flow/nodes/wren.py index 1f9e2c7..2048508 100644 --- a/backend/fluksio/flow/nodes/wren.py +++ b/backend/fluksio/flow/nodes/wren.py @@ -14,6 +14,7 @@ from collections.abc import Iterable from typing import Any from pydantic import BaseModel, ConfigDict, Field +from threadpoolctl import threadpool_limits from fluksio.flow.messages import DType, MessageSpec from fluksio.flow.nodes.base import Node @@ -306,7 +307,11 @@ class WrenNode(Node): wren = _wren() specs = {spec.port: spec for spec in self.output_ports} house = self._house(wren, specs, presence, outdoor, dark) - calls, confidence = wren.act(prompt, house) + # One BLAS thread answers in about half the time for a model this + # small. Scoped to the call, so numpy elsewhere in the process keeps + # its threads. + with threadpool_limits(1): + calls, confidence = wren.act(prompt, house) confidence = round(float(confidence), 3) write = self.cfg.apply and confidence >= self.cfg.min_confidence writes: dict[str, Any] = {} diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 18d7911..8e26983 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -58,6 +58,10 @@ dependencies = [ # The `mlp` node, which is registered at boot. Kept out of the extras # because a data-science environment has it already. "numpy>=2.2.6", + # The `wren` node runs its model on one BLAS thread, which answers in about + # half the time. Scoped to the call: OPENBLAS_NUM_THREADS would bind every + # other numpy call in the engine too, and only counts before numpy loads. + "threadpoolctl>=3.5", # The `inject` node's cron. Imported where it is used, behind a message # that cannot fire while this line is here. "croniter>=1.3.0", @@ -164,6 +168,11 @@ ignore_missing_imports = true module = ["wren"] ignore_missing_imports = true +# What scopes the `wren` node's BLAS threads to its call; it ships no types. +[[tool.mypy.overrides]] +module = ["threadpoolctl"] +ignore_missing_imports = true + [tool.ruff] target-version = "py312" diff --git a/backend/tests/flow/test_wren.py b/backend/tests/flow/test_wren.py index 2f3fff3..44c048b 100644 --- a/backend/tests/flow/test_wren.py +++ b/backend/tests/flow/test_wren.py @@ -143,3 +143,25 @@ def test_nothing_is_written_unless_applied_and_sure(fake_wren, params, confidenc } ] } + + +def test_the_model_runs_on_one_blas_thread(fake_wren, monkeypatch): + """About twice as fast for a model this small; numpy is loaded by `mlp`.""" + from threadpoolctl import threadpool_info + + threads: list[int] = [] + act = fake_wren.act + + def counted(command: str, house: Any) -> tuple[list[Any], float]: + threads.extend( + pool["num_threads"] + for pool in threadpool_info() + if pool["user_api"] == "blas" + ) + return act(command, house) + + monkeypatch.setattr(fake_wren, "act", counted) + _node(("kitchen_lamp", DType.BOOL)).execute({"house.command": "lights on"}) + + assert threads + assert set(threads) == {1} diff --git a/uv.lock b/uv.lock index 10e0f43..be2104a 100644 --- a/uv.lock +++ b/uv.lock @@ -946,6 +946,7 @@ dependencies = [ { name = "rich" }, { name = "sqlmodel" }, { name = "textual" }, + { name = "threadpoolctl" }, { name = "uv" }, { name = "websockets" }, ] @@ -998,6 +999,7 @@ requires-dist = [ { name = "sentry-sdk", extras = ["fastapi"], marker = "extra == 'server'", specifier = ">=2.20.0" }, { name = "sqlmodel", specifier = ">=0.0.21,<1.0.0" }, { name = "textual", specifier = ">=1.0" }, + { name = "threadpoolctl", specifier = ">=3.5" }, { name = "uv", specifier = ">=0.5" }, { name = "websockets", specifier = ">=14" }, ] @@ -2840,6 +2842,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/be/35261223d9416a0751cdff1c7b4a6f881387218a12d439fe22fefebc8c04/textual-8.2.8-py3-none-any.whl", hash = "sha256:267375fd402dc8d981457212efa71f0e3365fd17bba144ba9bb3ed7563cb374a", size = 731418, upload-time = "2026-06-30T06:51:26.364Z" }, ] +[[package]] +name = "threadpoolctl" +version = "3.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/00/dc/6c58154c1c65f758ea979e7139cb76993a9cfc662d14e9be3c4a667cfb77/threadpoolctl-3.7.0.tar.gz", hash = "sha256:61348cfb77d53b9242e0017029244b559b810c142ced65b4e21eeca1843959a7", size = 31961, upload-time = "2026-09-15T15:46:20.263Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/3f/f88a53f60a472b46f4023f56d204dd7de33d34c5d2acbfa0d70a674e639e/threadpoolctl-3.7.0-py3-none-any.whl", hash = "sha256:cd8b60b5641b45c67bbf73c64c843235fc2d8a480c87389f52f5dbee893b86be", size = 26362, upload-time = "2026-09-15T15:46:19.168Z" }, +] + [[package]] name = "typer" version = "0.21.1"