File size: 2,592 Bytes
9422383
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
"""Public Python runtime admission behavior that does not require Pi hardware."""

from __future__ import annotations

import hashlib
import json
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from arm_yolo26n import Yolo26nError, Yolo26nTensorSession


def _write_model(directory: Path, payload: bytes) -> None:
    directory.mkdir()
    (directory / "model.bin").write_bytes(payload)
    manifest = {
        "kind": "yolo26n-exported-model",
        "version": 1,
        "model": {
            "checkpoint_sha256": "0" * 64,
            "class_names": ["first", "second"],
            "input_shape": [1, 3, 320, 320],
            "output_shape": [1, 6, 2100],
            "target": {
                "version": 1,
                "target": "linux-aarch64-raspberry-pi-5",
                "input_shape": [1, 3, 320, 320],
                "input_dtype": "fp32",
                "output_dtype": "fp32",
                "thread_count": 4,
            },
        },
        "artifact": {
            "path": "model.bin",
            "bytes": len(payload),
            "sha256": hashlib.sha256(payload).hexdigest(),
        },
    }
    (directory / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")


class RuntimeAdmissionTests(unittest.TestCase):
    def test_session_requires_runtime_and_model_as_separate_inputs(self) -> None:
        with TemporaryDirectory() as directory:
            root = Path(directory)
            with self.assertRaisesRegex(FileNotFoundError, "runtime library"):
                Yolo26nTensorSession(runtime=root / "missing.so", model=root / "missing-model")

            runtime = root / "libarm-yolo26n.so.0"
            runtime.write_bytes(b"not loaded because model admission fails first")
            with self.assertRaisesRegex(FileNotFoundError, "model directory"):
                Yolo26nTensorSession(runtime=runtime, model=root / "missing-model")

    def test_session_rejects_a_model_that_does_not_match_its_manifest(self) -> None:
        with TemporaryDirectory() as directory:
            root = Path(directory)
            runtime = root / "libarm-yolo26n.so.0"
            runtime.write_bytes(b"not loaded because model admission fails first")
            model = root / "model"
            _write_model(model, b"model")
            (model / "model.bin").write_bytes(b"damaged")

            with self.assertRaisesRegex(Yolo26nError, "'model.bin' does not match its manifest"):
                Yolo26nTensorSession(runtime=runtime, model=model)


if __name__ == "__main__":
    unittest.main()