onnx--onnx/onnx/backend/test/runner/__init__.py

# Copyright (c) ONNX Project Contributors
#
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import functools
import glob
import os
import re
import shutil
import sys
import tempfile
import time
import unittest
from collections import defaultdict
from re import Pattern
from typing import TYPE_CHECKING, Any
from urllib.request import urlretrieve

import numpy as np

import onnx
import onnx.reference
from onnx import ONNX_ML, ModelProto, NodeProto, TypeProto, ValueInfoProto, numpy_helper
from onnx.backend.test.loader import load_model_tests
from onnx.backend.test.runner.item import TestItem

if TYPE_CHECKING:
    from collections.abc import Callable, Iterable, Sequence

    from onnx.backend.base import Backend
    from onnx.backend.test.case.test_case import TestCase


class BackendIsNotSupposedToImplementIt(unittest.SkipTest):
    pass


def retry_execute(times: int) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    assert times >= 1

    def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
        @functools.wraps(func)
        def wrapped(*args: Any, **kwargs: Any) -> Any:
            for i in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except Exception:  # noqa: PERF203
                    print(f"{i} times tried")
                    if i == times:
                        raise
                    time.sleep(5 * i)
            return None

        return wrapped

    return wrapper


class Runner:
    def __init__(
        self,
        backend: type[Backend],
        parent_module: str | None = None,
        test_kwargs: dict | None = None,
    ) -> None:
        self.backend = backend
        self._parent_module = parent_module
        self._include_patterns: set[Pattern[str]] = set()
        self._exclude_patterns: set[Pattern[str]] = set()
        self._xfail_patterns: set[Pattern[str]] = set()
        self._test_kwargs: dict = test_kwargs or {}

        # This is the source of the truth of all test functions.
        # Properties `test_cases`, `test_suite` and `tests` will be
        # derived from it.
        # {category: {name: func}}
        self._test_items: dict[str, dict[str, TestItem]] = defaultdict(dict)

        for rt in load_model_tests(kind="node"):
            self._add_model_test(rt, "Node")

        for rt in load_model_tests(kind="real"):
            self._add_model_test(rt, "Real")

        for rt in load_model_tests(kind="simple"):
            self._add_model_test(rt, "Simple")

        for ct in load_model_tests(kind="pytorch-converted"):
            self._add_model_test(ct, "PyTorchConverted")

        for test_case in load_model_tests(kind="pytorch-operator"):
            self._add_model_test(test_case, "PyTorchOperator")

    def _get_test_case(self, name: str) -> type[unittest.TestCase]:
        test_case = type(str(name), (unittest.TestCase,), {})
        if self._parent_module:
            test_case.__module__ = self._parent_module
        return test_case

    def include(self, pattern: str) -> Runner:
        self._include_patterns.add(re.compile(pattern))
        return self

    def exclude(self, pattern: str) -> Runner:
        self._exclude_patterns.add(re.compile(pattern))
        return self

    def xfail(self, pattern: str) -> Runner:
        self._xfail_patterns.add(re.compile(pattern))
        return self

    def enable_report(self) -> Runner:
        import pytest  # noqa: PLC0415

        for category, items_map in self._test_items.items():
            for item in items_map.values():
                item.func = pytest.mark.onnx_coverage(item.proto, category)(item.func)
        return self

    @property
    def _filtered_test_items(self) -> dict[str, dict[str, TestItem]]:
        filtered: dict[str, dict[str, TestItem]] = {}
        for category, items_map in self._test_items.items():
            filtered[category] = {}
            for name, item in items_map.items():
                if self._include_patterns and (
                    not any(include.search(name) for include in self._include_patterns)
                ):
                    item.func = unittest.skip("no matched include pattern")(item.func)
                for exclude in self._exclude_patterns:
                    if exclude.search(name):
                        item.func = unittest.skip(
                            f'matched exclude pattern "{exclude.pattern}"'
                        )(item.func)
                for xfail in self._xfail_patterns:
                    if xfail.search(name):
                        item.func = unittest.expectedFailure(item.func)
                filtered[category][name] = item
        return filtered

    @property
    def test_cases(self) -> dict[str, type[unittest.TestCase]]:
        """List of test cases to be applied on the parent scope
        Example usage:
            globals().update(BackendTest(backend).test_cases)
        """
        test_cases = {}
        for category, items_map in self._filtered_test_items.items():
            test_case_name = f"OnnxBackend{category}Test"
            test_case = self._get_test_case(test_case_name)
            for name, item in sorted(items_map.items()):
                setattr(test_case, name, item.func)
            test_cases[test_case_name] = test_case
        return test_cases

    @property
    def test_suite(self) -> unittest.TestSuite:
        """TestSuite that can be run by TestRunner
        Example usage:
            unittest.TextTestRunner().run(BackendTest(backend).test_suite)
        """
        suite = unittest.TestSuite()
        for case in sorted(
            self.test_cases.values(), key=lambda cl: cl.__class__.__name__
        ):
            suite.addTests(unittest.defaultTestLoader.loadTestsFromTestCase(case))
        return suite

    # For backward compatibility (we used to expose `.tests`)
    @property
    def tests(self) -> type[unittest.TestCase]:
        """One single unittest.TestCase that hosts all the test functions
        Example usage:
            onnx_backend_tests = BackendTest(backend).tests
        """
        tests = self._get_test_case("OnnxBackendTest")
        for items_map in sorted(
            self._filtered_test_items.values(), key=lambda cl: cl.__class__.__name__
        ):
            for name, item in sorted(items_map.items()):
                setattr(tests, name, item.func)
        return tests

    @classmethod
    def assert_similar_outputs(
        cls,
        ref_outputs: Sequence[Any],
        outputs: Sequence[Any],
        rtol: float,
        atol: float,
        model_dir: str | None = None,
    ) -> None:
        try:
            np.testing.assert_equal(len(outputs), len(ref_outputs))
        except TypeError as e:
            raise TypeError(
                f"Unable to compare expected type {type(ref_outputs)} "
                f"and runtime type {type(outputs)} (known test={model_dir or '?'!r})"
            ) from e
        for i in range(len(outputs)):
            if isinstance(outputs[i], (list, tuple)):
                if not isinstance(ref_outputs[i], (list, tuple)):
                    raise AssertionError(  # noqa: TRY004
                        f"Unexpected type {type(outputs[i])} for outputs[{i}]. Expected "
                        f"type is {type(ref_outputs[i])} (known test={model_dir or '?'!r})."
                    )
                for j in range(len(outputs[i])):
                    cls.assert_similar_outputs(
                        ref_outputs[i][j],
                        outputs[i][j],
                        rtol,
                        atol,
                        model_dir=model_dir,
                    )
            else:
                np.testing.assert_array_equal(
                    outputs[i].shape,
                    ref_outputs[i].shape,
                    err_msg=f"Output {i} has incorrect shape",
                )
                if ref_outputs[i].dtype == object:
                    # Strings
                    np.testing.assert_array_equal(outputs[i], ref_outputs[i])
                else:
                    np.testing.assert_equal(outputs[i].dtype, ref_outputs[i].dtype)
                    np.testing.assert_allclose(
                        outputs[i], ref_outputs[i], rtol=rtol, atol=atol
                    )

    @classmethod
    @retry_execute(3)
    def download_model(
        cls,
        model_test: TestCase,
        models_dir: str,
    ) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            try:
                assert model_test.url
                print(
                    f"Start downloading model {model_test.model_name} from {model_test.url}"
                )
                filename = os.path.join(tmpdir, "file")
                urlretrieve(model_test.url, filename)
                print("Done")
                onnx.utils._extract_model_safe(filename, models_dir)
            except Exception as e:
                print(f"Failed to prepare data for model {model_test.model_name}: {e}")
                raise

    @classmethod
    def prepare_model_data(cls, model_test: TestCase) -> str:
        onnx_home = os.path.expanduser(
            os.getenv("ONNX_HOME", os.path.join("~", ".onnx"))
        )
        models_dir = os.getenv("ONNX_MODELS", os.path.join(onnx_home, "models"))
        model_dir: str = os.path.join(models_dir, model_test.model_name)
        if not os.path.exists(os.path.join(model_dir, "model.onnx")):
            if os.path.exists(model_dir):
                bi = 0
                while True:
                    dest = f"{model_dir}.old.{bi}"
                    if os.path.exists(dest):
                        bi += 1
                        continue
                    shutil.move(model_dir, dest)
                    break
            os.makedirs(model_dir)

            cls.download_model(model_test=model_test, models_dir=models_dir)
        return model_dir

    def _add_test(
        self,
        category: str,
        test_name: str,
        test_func: Callable[..., Any],
        report_item: list[ModelProto | NodeProto | None],
        devices: Iterable[str] = ("CPU", "CUDA"),
        **kwargs: Any,
    ) -> None:
        # We don't prepend the 'test_' prefix to improve greppability
        if not test_name.startswith("test_"):
            raise ValueError(f"Test name must start with test_: {test_name}")

        def add_device_test(device: str) -> None:
            device_test_name = f"{test_name}_{device.lower()}"
            if device_test_name in self._test_items[category]:
                raise ValueError(
                    f'Duplicated test name "{device_test_name}" in category "{category}"'
                )

            @unittest.skipIf(
                not self.backend.supports_device(device),
                f"Backend doesn't support device {device}",
            )
            @functools.wraps(test_func)
            def device_test_func(*args: Any, **device_test_kwarg: Any) -> Any:
                try:
                    merged_kwargs = {**kwargs, **device_test_kwarg}
                    return test_func(*args, device, **merged_kwargs)
                except BackendIsNotSupposedToImplementIt as e:
                    # hacky verbose reporting
                    if "-v" in sys.argv or "--verbose" in sys.argv:
                        print(f"Test {device_test_name} is effectively skipped: {e}")

            self._test_items[category][device_test_name] = TestItem(
                device_test_func, report_item
            )

        for device in devices:
            add_device_test(device)

    @staticmethod
    def generate_dummy_data(
        x: ValueInfoProto, seed: int = 0, name: str = "", random: bool = False
    ) -> np.ndarray:
        """Generates a random tensor based on the input definition."""
        if not x.type.tensor_type:
            raise NotImplementedError(
                f"Input expected to have tensor type. "
                f"Unable to generate random data for model {name!r} and input {x}."
            )
        if x.type.tensor_type.elem_type != 1:
            raise NotImplementedError(
                f"Currently limited to float tensors. "
                f"Unable to generate random data for model {name!r} and input {x}."
            )
        shape = tuple(
            d.dim_value if d.HasField("dim_value") else 1
            for d in x.type.tensor_type.shape.dim
        )
        if random:
            gen = np.random.default_rng(seed=seed)
            return gen.random(shape, np.float32)
        n = np.prod(shape)
        return (np.arange(n).reshape(shape) / n).astype(np.float32)

    def _add_model_test(self, model_test: TestCase, kind: str) -> None:
        # model is loaded at runtime, note sometimes it could even
        # never loaded if the test skipped
        model_marker: list[ModelProto | NodeProto | None] = [None]

        def run(test_self: Any, device: str, **kwargs) -> None:  # noqa: ARG001
            if model_test.url is not None and model_test.url.startswith(
                "onnx/backend/test/data/light/"
            ):
                # testing local files
                model_pb_path = os.path.normpath(
                    os.path.join(
                        os.path.dirname(__file__),
                        "..",
                        "..",
                        "..",
                        "..",
                        model_test.url,
                    )
                )
                if not os.path.exists(model_pb_path):
                    raise FileNotFoundError(f"Unable to find model {model_pb_path!r}.")
                onnx_home = os.path.expanduser(
                    os.getenv("ONNX_HOME", os.path.join("~", ".onnx"))
                )
                models_dir = os.getenv(
                    "ONNX_MODELS", os.path.join(onnx_home, "models", "light")
                )
                model_dir: str = os.path.join(models_dir, model_test.model_name)
                if not os.path.exists(model_dir):
                    os.makedirs(model_dir)
                use_dummy = True
            else:
                if model_test.model_dir is None:
                    model_dir = self.prepare_model_data(model_test)
                else:
                    model_dir = model_test.model_dir
                model_pb_path = os.path.join(model_dir, "model.onnx")
                use_dummy = False

            if not ONNX_ML and "ai_onnx_ml" in model_dir:
                return

            model = onnx.load(model_pb_path)
            model_marker[0] = model
            if (
                hasattr(self.backend, "is_compatible")
                and callable(self.backend.is_compatible)
                and not self.backend.is_compatible(model)
            ):
                raise unittest.SkipTest("Not compatible with backend")

            prepared_model = self.backend.prepare(model, device, **kwargs)
            assert prepared_model is not None

            if use_dummy:
                # When the backend test goes through a test involving a
                # model stored in onnx/backend/test/data/light,
                # this function generates expected output coming from
                # from ReferenceEvaluator run with random inputs.
                # A couple of models include many Conv operators and the
                # python implementation is slow (such as test_bvlc_alexnet).
                with open(model_pb_path, "rb") as f:
                    onx = onnx.load(f)

                test_data_set = os.path.join(model_dir, "test_data_set_0")
                if not os.path.exists(test_data_set):
                    os.mkdir(test_data_set)
                feeds = {}
                inits = {i.name for i in onx.graph.initializer}
                n_input = 0
                inputs = []
                for i in range(len(onx.graph.input)):
                    if onx.graph.input[i].name in inits:
                        continue
                    name = os.path.join(test_data_set, f"input_{n_input}.pb")
                    inputs.append(name)
                    n_input += 1
                    x = onx.graph.input[i]
                    value = self.generate_dummy_data(
                        x, seed=0, name=model_test.model_name, random=False
                    )
                    feeds[x.name] = value
                    with open(name, "wb") as f:
                        f.write(onnx.numpy_helper.from_array(value).SerializeToString())

                # loads expected output if any available
                prefix = os.path.splitext(model_pb_path)[0]
                expected_outputs = []
                for i in range(len(onx.graph.output)):
                    name = f"{prefix}_output_{i}.pb"
                    if os.path.exists(name):
                        expected_outputs.append(name)
                        continue
                    expected_outputs = None
                    break

                if expected_outputs is None:
                    ref = onnx.reference.ReferenceEvaluator(onx)
                    outputs = ref.run(None, feeds)
                    for i, o in enumerate(outputs):
                        name = os.path.join(test_data_set, f"output_{i}.pb")
                        with open(name, "wb") as f:
                            f.write(onnx.numpy_helper.from_array(o).SerializeToString())
                else:
                    for i, o in enumerate(expected_outputs):
                        name = os.path.join(test_data_set, f"output_{i}.pb")
                        shutil.copy(o, name)
            else:
                # TODO after converting all npz files to protobuf, we can delete this.
                for test_data_npz in glob.glob(
                    os.path.join(model_dir, "test_data_*.npz")
                ):
                    test_data = np.load(test_data_npz, encoding="bytes")
                    inputs = list(test_data["inputs"])
                    outputs = list(prepared_model.run(inputs))
                    ref_outputs = tuple(
                        np.array(x) if not isinstance(x, (list, dict)) else x
                        for f in test_data["outputs"]
                    )
                    self.assert_similar_outputs(
                        ref_outputs,
                        outputs,
                        rtol=kwargs.get("rtol", model_test.rtol),
                        atol=kwargs.get("atol", model_test.atol),
                        model_dir=model_dir,
                    )

            for test_data_dir in glob.glob(os.path.join(model_dir, "test_data_set*")):
                inputs = []
                inputs_num = len(glob.glob(os.path.join(test_data_dir, "input_*.pb")))
                for i in range(inputs_num):
                    input_file = os.path.join(test_data_dir, f"input_{i}.pb")
                    self._load_proto(input_file, inputs, model.graph.input[i].type)
                ref_outputs = []
                ref_outputs_num = len(
                    glob.glob(os.path.join(test_data_dir, "output_*.pb"))
                )
                for i in range(ref_outputs_num):
                    output_file = os.path.join(test_data_dir, f"output_{i}.pb")
                    self._load_proto(
                        output_file, ref_outputs, model.graph.output[i].type
                    )
                outputs = list(prepared_model.run(inputs))
                self.assert_similar_outputs(
                    ref_outputs,
                    outputs,
                    rtol=kwargs.get("rtol", model_test.rtol),
                    atol=kwargs.get("atol", model_test.atol),
                    model_dir=model_dir,
                )

        if model_test.name in self._test_kwargs:
            self._add_test(
                kind + "Model",
                model_test.name,
                run,
                model_marker,
                **self._test_kwargs[model_test.name],
            )
        else:
            self._add_test(kind + "Model", model_test.name, run, model_marker)

    def _load_proto(
        self,
        proto_filename: str,
        target_list: list[np.ndarray | list[Any]],
        model_type_proto: TypeProto,
    ) -> None:
        with open(proto_filename, "rb") as f:
            protobuf_content = f.read()
            if model_type_proto.HasField("sequence_type"):
                sequence = onnx.SequenceProto()
                sequence.ParseFromString(protobuf_content)
                target_list.append(numpy_helper.to_list(sequence))
            elif model_type_proto.HasField("tensor_type"):
                tensor = onnx.TensorProto()
                tensor.ParseFromString(protobuf_content)
                t = numpy_helper.to_array(tensor)
                assert isinstance(t, np.ndarray)
                target_list.append(t)
            elif model_type_proto.HasField("optional_type"):
                optional = onnx.OptionalProto()
                optional.ParseFromString(protobuf_content)
                target_list.append(numpy_helper.to_optional(optional))  # type: ignore[arg-type]
            else:
                print(
                    "Loading proto of that specific type (Map/Sparse Tensor) is currently not supported"
                )