From e3761617443440c8cb5471b6531c320f98e82906 Mon Sep 17 00:00:00 2001 From: Jinman Park <67285913+parkjinman98@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:15:43 +0900 Subject: [PATCH 1/2] fix: bound multi-slot backend allocations --- .agents/skills/mblt-npu/SKILL.md | 1 + AGENTS.md | 5 ++-- README.md | 4 +++ mblt_npu/__init__.py | 4 +++ mblt_npu/npu_backend.py | 40 +++++++++++++++++++++++----- tests/test_npu_backend_multi_slot.py | 39 +++++++++++++++++++++++++++ 6 files changed, 84 insertions(+), 9 deletions(-) diff --git a/.agents/skills/mblt-npu/SKILL.md b/.agents/skills/mblt-npu/SKILL.md index a36cf24..be52942 100644 --- a/.agents/skills/mblt-npu/SKILL.md +++ b/.agents/skills/mblt-npu/SKILL.md @@ -14,6 +14,7 @@ description: >- 2. Preserve the `MobilintNPUBackend` compatibility contract: lifecycle methods, callable MXQ inference, input dtype inspection, serialization, and board-selected subclasses. For multi-slot backends, `max_batch_size` is aggregate capacity; keep slot-zero handles writable, + accept only non-boolean integers from 1 through 1024, cap each backend at 64 model slots, preserve `infer_slot()`, and roll back every slot after allocation or launch failure. 3. Normalize legacy `aries`/`regulus` target values to board identifiers. Accepted boards are `aries-rb`, `regulus-ra`, `regulus-rb`, `regulus-ra-usb`, and `regulus-rb-usb`. Forward the diff --git a/AGENTS.md b/AGENTS.md index dcc17c6..aa2bcab 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -26,8 +26,9 @@ aligned with that floor before using newer language syntax or standard-library A selection behavior. - `max_batch_size` is aggregate capacity: probe the compiled per-slot batch axis `K`, launch `ceil(max_batch_size / K)` slots, share accelerators per device, - and preserve slot-zero compatibility handles. Keep allocation rollback and - `infer_slot()` deterministic across all slots. + and preserve slot-zero compatibility handles. Accept only non-boolean integers + from 1 through 1024 and cap the derived allocation at 64 model slots. Keep + allocation rollback and `infer_slot()` deterministic across all slots. - Normalize legacy `aries` to `aries-rb` and `regulus` to `regulus-ra`. Supported board identifiers are `aries-rb`, `regulus-ra`, `regulus-rb`, `regulus-ra-usb`, and `regulus-rb-usb`. The board name is passed as the first argument to diff --git a/README.md b/README.md index e68d232..f28e7ae 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,10 @@ need to inspect the supported targets. `max_batch_size` is aggregate capacity. At `create()`, the backend probes the compiled per-model capacity `K` and loads `ceil(max_batch_size / K)` model slots. +For predictable resource use, `max_batch_size` must be a non-boolean integer +from 1 through 1024, and one backend may load at most 64 model slots. If the +probed `K` would require more slots, creation disposes slot zero and fails before +allocating any additional slots. Slots are distributed round-robin over the devices named by canonical target strings and reuse one accelerator per device. `mxq_model` and `acc` continue to refer to slot zero for compatibility; concurrent callers can use diff --git a/mblt_npu/__init__.py b/mblt_npu/__init__.py index 4fb4f68..744acd7 100644 --- a/mblt_npu/__init__.py +++ b/mblt_npu/__init__.py @@ -12,6 +12,8 @@ from .npu_backend import ( BACKEND_CLASSES, DEFAULT_TARGET_DEVICE, + MAX_BATCH_SIZE, + MAX_MODEL_SLOTS, MobilintBackendAllocError, MobilintAriesBackend, MobilintNPUBackend, @@ -27,6 +29,8 @@ __all__ = [ "BACKEND_CLASSES", "DEFAULT_TARGET_DEVICE", + "MAX_BATCH_SIZE", + "MAX_MODEL_SLOTS", "MobilintBackendAllocError", "MobilintAriesBackend", "MobilintNPUBackend", diff --git a/mblt_npu/npu_backend.py b/mblt_npu/npu_backend.py index 9efa0c2..5320724 100644 --- a/mblt_npu/npu_backend.py +++ b/mblt_npu/npu_backend.py @@ -27,7 +27,6 @@ """ import logging -import math import os import re import sys @@ -62,6 +61,12 @@ DEFAULT_TARGET_DEVICE = "aries-rb" """Default supported Mobilint NPU board.""" +MAX_BATCH_SIZE = 1024 +"""Largest aggregate batch capacity accepted by the public configuration API.""" + +MAX_MODEL_SLOTS = 64 +"""Largest number of native model slots a backend instance may allocate.""" + _TARGET_DEVICE_ALIASES = {"aries": "aries-rb", "regulus": "regulus-ra"} @@ -242,9 +247,12 @@ def __init__( expanded into ``target_cores`` / ``target_clusters`` when those lists are empty, and prepends the device prefix to legacy 2-part items. - max_batch_size: Requested aggregate batch capacity. The backend + max_batch_size: Requested aggregate batch capacity in the inclusive + range ``1..MAX_BATCH_SIZE``. Booleans and non-integers are rejected. + The backend launches enough slots so that ``N * K >= max_batch_size``, - where ``K`` is the compiled batch axis of the MXQ artifact. + where ``K`` is the compiled batch axis of the MXQ artifact, up + to ``MAX_MODEL_SLOTS`` slots. core_mode: Execution mode that determines how NPU cores are allocated. One of ``"single"``, ``"multi"``, ``"global4"``, or ``"global8"``. @@ -278,7 +286,14 @@ def __init__( self.revision = revision self._commit_hash = commit_hash self.mxq_path = mxq_path - self.max_batch_size = max(1, max_batch_size) + if isinstance(max_batch_size, bool) or not isinstance(max_batch_size, int): + raise TypeError("max_batch_size must be a non-boolean integer.") + if not 1 <= max_batch_size <= MAX_BATCH_SIZE: + raise ValueError( + f"max_batch_size must be between 1 and {MAX_BATCH_SIZE}, " + f"got {max_batch_size}." + ) + self.max_batch_size = max_batch_size # Multi-slot backing state; populated in create()/launch(). # ``self.acc`` and ``self.mxq_model`` remain accessible as @@ -1053,9 +1068,20 @@ def _spawn_slot(slot_idx: int, dev: int) -> "Model": file=sys.stderr, ) raise - self.n_models = max( - 1, math.ceil(self.max_batch_size / max(1, self.k_per_model)) - ) + # Integer ceiling avoids converting attacker-controlled integers to a + # float. Keep this post-probe guard even though construction validates + # max_batch_size: the attribute remains writable for compatibility. + k_per_model = max(1, self.k_per_model) + self.n_models = (self.max_batch_size + k_per_model - 1) // k_per_model + if self.n_models > MAX_MODEL_SLOTS: + self._dispose_all_slots() + self.accs = {} + self.n_models = 0 + raise ValueError( + f"max_batch_size={self.max_batch_size} with " + f"k_per_model={k_per_model} requires more than the supported " + f"maximum of {MAX_MODEL_SLOTS} model slots." + ) for slot_idx in range(1, self.n_models): d = int(unique_devs[slot_idx % len(unique_devs)]) diff --git a/tests/test_npu_backend_multi_slot.py b/tests/test_npu_backend_multi_slot.py index e29adf4..98a4586 100644 --- a/tests/test_npu_backend_multi_slot.py +++ b/tests/test_npu_backend_multi_slot.py @@ -21,6 +21,8 @@ from mblt_npu import npu_backend as npu_backend_module from mblt_npu.npu_backend import ( + MAX_BATCH_SIZE, + MAX_MODEL_SLOTS, MobilintBackendAllocError, MobilintNPUBackend, _is_qbruntime_bad_alloc, @@ -153,6 +155,43 @@ def _make_backend_at(tmp_path, **kwargs) -> MobilintNPUBackend: return MobilintNPUBackend(**kwargs) +@pytest.mark.parametrize("value", [True, False, 1.5, "2", None]) +def test_max_batch_size_rejects_non_integer_values(value) -> None: + """Public construction rejects values that are not strict integers.""" + with pytest.raises(TypeError, match="non-boolean integer"): + MobilintNPUBackend(max_batch_size=value) + + +@pytest.mark.parametrize("value", [0, -1, MAX_BATCH_SIZE + 1]) +def test_max_batch_size_rejects_values_outside_safe_range(value: int) -> None: + """Public construction bounds aggregate capacity before native allocation.""" + with pytest.raises(ValueError, match=f"between 1 and {MAX_BATCH_SIZE}"): + MobilintNPUBackend(max_batch_size=value) + + +def test_from_dict_validates_max_batch_size() -> None: + """Deserialized configuration uses the same strict admission checks.""" + with pytest.raises(TypeError, match="non-boolean integer"): + MobilintNPUBackend.from_dict({"max_batch_size": 2.5}) + + +def test_create_rejects_excessive_derived_slots_and_rolls_back( + tmp_path, stub_qbruntime +) -> None: + """The post-probe slot cap disposes slot zero before further allocations.""" + backend = _make_backend_at(tmp_path, max_batch_size=MAX_MODEL_SLOTS) + backend.max_batch_size = MAX_MODEL_SLOTS + 1 + + with pytest.raises(ValueError, match=f"maximum of {MAX_MODEL_SLOTS} model slots"): + backend.create() + + assert len(stub_qbruntime.models) == 1 + assert stub_qbruntime.models[0].disposed is True + assert backend.mxq_models == [] + assert backend.accs == {} + assert backend.n_models == 0 + + def test_backend_create_n1_single_slot_single_device(tmp_path, stub_qbruntime) -> None: """max_batch_size=1 with a K=1 model yields exactly one slot on the target device.""" backend = _make_backend_at( From d152208b390c566e360e77260fb3ff0d3a4c6705 Mon Sep 17 00:00:00 2001 From: jinman Date: Tue, 22 Sep 2026 17:55:08 +0900 Subject: [PATCH 2/2] fix: revalidate batch size before allocation --- mblt_npu/npu_backend.py | 31 +++++++++++++++++++--------- tests/test_npu_backend_multi_slot.py | 22 ++++++++++++++++++++ 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/mblt_npu/npu_backend.py b/mblt_npu/npu_backend.py index 5320724..0d6dcc9 100644 --- a/mblt_npu/npu_backend.py +++ b/mblt_npu/npu_backend.py @@ -70,6 +70,17 @@ _TARGET_DEVICE_ALIASES = {"aries": "aries-rb", "regulus": "regulus-ra"} +def _validate_max_batch_size(value: object) -> int: + """Return a safe aggregate capacity or reject it before native allocation.""" + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError("max_batch_size must be a non-boolean integer.") + if not 1 <= value <= MAX_BATCH_SIZE: + raise ValueError( + f"max_batch_size must be between 1 and {MAX_BATCH_SIZE}, got {value}." + ) + return value + + def normalize_target_device(target_device: str) -> str: """Normalize a board identifier while accepting legacy product aliases.""" if not isinstance(target_device, str): @@ -286,14 +297,7 @@ def __init__( self.revision = revision self._commit_hash = commit_hash self.mxq_path = mxq_path - if isinstance(max_batch_size, bool) or not isinstance(max_batch_size, int): - raise TypeError("max_batch_size must be a non-boolean integer.") - if not 1 <= max_batch_size <= MAX_BATCH_SIZE: - raise ValueError( - f"max_batch_size must be between 1 and {MAX_BATCH_SIZE}, " - f"got {max_batch_size}." - ) - self.max_batch_size = max_batch_size + self.max_batch_size = _validate_max_batch_size(max_batch_size) # Multi-slot backing state; populated in create()/launch(). # ``self.acc`` and ``self.mxq_model`` remain accessible as @@ -966,15 +970,22 @@ def create(self) -> None: user-config or artifact bug. Raises: + TypeError: If the current ``max_batch_size`` is not a non-boolean + integer. MobilintBackendAllocError: If any slot hits a device-memory ``BadAlloc`` or the slot 0 K probe hits a ``BadAlloc``. QbRuntimeError: If any slot or the slot 0 K probe fails for a non-alloc reason (after partial-state rollback). - ValueError: If ``self.core_mode`` is not one of the supported - values. + ValueError: If the current ``max_batch_size`` is outside the + supported range or ``self.core_mode`` is unsupported. AssertionError: If ``"global8"`` mode is requested but a device does not cover both clusters. """ + # ``max_batch_size`` remains writable for compatibility. Revalidate + # its current value before opening accelerators or creating slot 0 so + # post-construction mutation cannot bypass the public capacity bounds. + self.max_batch_size = _validate_max_batch_size(self.max_batch_size) + unique_devs = self._unique_devs_from_targets() if not unique_devs: unique_devs = [self._fallback_dev()] diff --git a/tests/test_npu_backend_multi_slot.py b/tests/test_npu_backend_multi_slot.py index 98a4586..e78b7d7 100644 --- a/tests/test_npu_backend_multi_slot.py +++ b/tests/test_npu_backend_multi_slot.py @@ -175,6 +175,28 @@ def test_from_dict_validates_max_batch_size() -> None: MobilintNPUBackend.from_dict({"max_batch_size": 2.5}) +@pytest.mark.parametrize( + ("value", "error", "message"), + [ + (True, TypeError, "non-boolean integer"), + (0, ValueError, f"between 1 and {MAX_BATCH_SIZE}"), + (MAX_BATCH_SIZE + 1, ValueError, f"between 1 and {MAX_BATCH_SIZE}"), + ], +) +def test_max_batch_size_revalidates_writable_assignments( + tmp_path, stub_qbruntime, value, error, message +) -> None: + """Compatibility writes are revalidated before any native allocation.""" + backend = _make_backend_at(tmp_path, max_batch_size=1) + backend.max_batch_size = value + + with pytest.raises(error, match=message): + backend.create() + + assert stub_qbruntime.models == [] + assert backend.accs == {} + + def test_create_rejects_excessive_derived_slots_and_rolls_back( tmp_path, stub_qbruntime ) -> None: