chore: import upstream snapshot with attribution

tools/Polygraphy/tests/util/test_util.py

@@ -0,0 +1,420 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import contextlib
+import io
+import os
+import random
+import tempfile
+from multiprocessing import Process
+
+import numpy as np
+import pytest
+
+from polygraphy import util
+from polygraphy.backend.trt import engine_from_network, network_from_onnx_bytes
+from polygraphy.util import util as util_internal  # For accessing and testing private functions in util.py
+from tests.models.meta import ONNX_MODELS
+
+VOLUME_CASES = [
+    ((1, 1, 1), 1),
+    ((2, 3, 4), 24),
+    (tuple(), 1),
+]
+
+
+@pytest.mark.parametrize("case", VOLUME_CASES)
+def test_volume(case):
+    it, vol = case
+    assert util.volume(it) == vol
+
+
+class FindStrInIterableCase:
+    def __init__(self, name, seq, index, expected):
+        self.name = name
+        self.seq = seq
+        self.index = index
+        self.expected = expected
+
+
+FIND_STR_IN_ITERABLE_CASES = [
+    # Case insensitve, plus function should return element from sequence, not name.
+    FindStrInIterableCase(
+        "Softmax:0", seq=["Softmax:0"], index=None, expected="Softmax:0"
+    ),
+    FindStrInIterableCase(
+        "Softmax:0", seq=["softmax:0"], index=None, expected="softmax:0"
+    ),
+    # Exact matches should take priority
+    FindStrInIterableCase(
+        "exact_name",
+        seq=["exact_name_plus", "exact_name"],
+        index=0,
+        expected="exact_name",
+    ),
+    # Index should come into play when no matches are found
+    FindStrInIterableCase(
+        "non-existent", seq=["test", "test2"], index=1, expected="test2"
+    ),
+]
+
+
+@pytest.mark.parametrize("case", FIND_STR_IN_ITERABLE_CASES)
+def test_find_str_in_iterable(case):
+    actual = util.find_str_in_iterable(case.name, case.seq, case.index)
+    assert actual == case.expected
+
+
+SHAPE_OVERRIDE_CASES = [
+    ((1, 3, 224, 224), (None, 3, 224, 224), True),
+]
+
+
+@pytest.mark.parametrize("case", SHAPE_OVERRIDE_CASES)
+def test_is_valid_shape_override(case):
+    override, shape, expected = case
+    assert (
+        util.is_valid_shape_override(new_shape=override, original_shape=shape)
+        == expected
+    )
+
+
+def arange(shape):
+    return np.arange(util.volume(shape)).reshape(shape)
+
+
+SHAPE_MATCHING_CASES = [
+    (arange((1, 1, 3, 3)), (3, 3), arange((3, 3))),  # Squeeze array shape
+    (
+        arange((1, 3, 3, 1)),
+        (1, 1, 3, 3),
+        arange((1, 1, 3, 3)),
+    ),  # Permutation should make no difference as other dimensions are 1s
+    (arange((3, 3)), (1, 1, 3, 3), arange((1, 1, 3, 3))),  # Unsqueeze where needed
+    (arange((3, 3)), (-1, 3), arange((3, 3))),  # Infer dynamic
+    (
+        arange((3 * 2 * 2,)),
+        (None, 3, 2, 2),
+        arange((1, 3, 2, 2)),
+    ),  # Reshape with inferred dimension
+    (
+        arange((1, 3, 2, 2)),
+        (None, 2, 2, 3),
+        np.transpose(arange((1, 3, 2, 2)), [0, 2, 3, 1]),
+    ),  # Permute
+]
+
+build_torch = lambda a, **kwargs: util.array.to_torch(np.array(a, **kwargs))
+
+
+@pytest.mark.parametrize("array_type", [np.array, build_torch])
+@pytest.mark.parametrize("arr, shape, expected", SHAPE_MATCHING_CASES)
+def test_shape_matching(arr, shape, expected, array_type):
+    arr = util.try_match_shape(array_type(arr), shape)
+    assert util.array.equal(arr, array_type(expected))
+
+
+UNPACK_ARGS_CASES = [
+    ((0, 1, 2), 3, (0, 1, 2)),  # no extras
+    ((0, 1, 2), 4, (0, 1, 2, None)),  # 1 extra
+    ((0, 1, 2), 2, (0, 1)),  # 1 fewer
+]
+
+
+@pytest.mark.parametrize("case", UNPACK_ARGS_CASES)
+def test_unpack_args(case):
+    args, num, expected = case
+    assert util.unpack_args(args, num) == expected
+
+
+UNIQUE_LIST_CASES = [
+    ([], []),
+    ([3, 1, 2], [3, 1, 2]),
+    ([1, 2, 3, 2, 1], [1, 2, 3]),
+    ([0, 0, 0, 0, 1, 0, 0], [0, 1]),
+    ([5, 5, 5, 5, 5], [5]),
+]
+
+
+@pytest.mark.parametrize("case", UNIQUE_LIST_CASES)
+def test_unique_list(case):
+    lst, expected = case
+    assert util.unique_list(lst) == expected
+
+
+def test_find_in_dirs():
+    with tempfile.TemporaryDirectory() as topdir:
+        dirs = list(
+            map(
+                lambda x: os.path.join(topdir, x),
+                ["test0", "test1", "test2", "test3", "test4"],
+            )
+        )
+        for subdir in dirs:
+            os.makedirs(subdir)
+
+        path_dir = random.choice(dirs)
+        path = os.path.join(path_dir, "cudart64_11.dll")
+
+        with open(path, "w") as f:
+            f.write("This file should be found by find_in_dirs")
+
+        assert util.find_in_dirs("cudart64_*.dll", dirs) == [path]
+
+
+@pytest.mark.parametrize(
+    "val,key,default,expected",
+    [
+        (1.0, None, None, 1.0),  # Basic
+        ({"inp": "hi"}, "inp", "", "hi"),  # Per-key
+        ({"inp": "hi"}, "out", "default", "default"),  # Per-key missing
+        ({"inp": 1.0, "": 2.0}, "out", 1.5, 2.0),  # Per-key with default
+    ],
+)
+def test_value_or_from_dict(val, key, default, expected):
+    actual = util.value_or_from_dict(val, key, default)
+    assert actual == expected
+
+
+def test_atomic_open():
+    def write_to_file(path, content):
+        with util.LockFile(path):
+            old_contents = util.load_file(path, mode="r")
+            util.save_file(old_contents + content, path, mode="w")
+
+    NUM_LINES = 10
+    NUM_PROCESSES = 5
+
+    outfile = util.NamedTemporaryFile()
+
+    processes = [
+        Process(
+            target=write_to_file,
+            args=(outfile.name, f"{proc} - writing line\n" * NUM_LINES),
+        )
+        for proc in range(NUM_PROCESSES)
+    ]
+
+    for process in processes:
+        process.start()
+
+    for process in processes:
+        process.join()
+
+    for process in processes:
+        assert not process.is_alive()
+        assert process.exitcode == 0
+
+    # Since we write atomically, all processes should be able to write their
+    # contents. Furthermore, the contents should be grouped by process.
+    with open(outfile.name) as f:
+        lines = list(f.readlines())
+        assert len(lines) == NUM_LINES * NUM_PROCESSES
+
+        for idx in range(NUM_PROCESSES):
+            offset = idx * NUM_LINES
+            expected_prefix = lines[offset].partition("-")[0].strip()
+            assert all(
+                line.startswith(expected_prefix)
+                for line in lines[offset : offset + NUM_LINES]
+            )
+
+    # Make sure the lock file is written to the correct path and not removed automatically.
+    assert os.path.exists(outfile.name + ".lock")
+
+
+class TestMakeRepr:
+    def test_basic(self):
+        assert util.make_repr("Example", 1, x=2) == ("Example(1, x=2)", False, False)
+
+    def test_default_args(self):
+        assert util.make_repr("Example", None, None, x=2) == (
+            "Example(None, None, x=2)",
+            True,
+            False,
+        )
+
+    def test_empty_args_are_default(self):
+        assert util.make_repr("Example", x=2) == ("Example(x=2)", True, False)
+
+    def test_default_kwargs(self):
+        assert util.make_repr("Example", 1, 2, x=None, y=None) == (
+            "Example(1, 2)",
+            False,
+            True,
+        )
+
+    def test_empty_kwargs_are_default(self):
+        assert util.make_repr("Example", 1, 2) == ("Example(1, 2)", False, True)
+
+    def test_does_not_modify(self):
+        obj = {"x": float("inf")}
+        assert util.make_repr("Example", obj) == (
+            "Example({'x': float('inf')})",
+            False,
+            True,
+        )
+        assert obj == {"x": float("inf")}
+
+    @pytest.mark.parametrize("obj", [float("nan"), float("inf"), float("-inf")])
+    @pytest.mark.parametrize("recursion_depth", [0, 1, 2])
+    def test_nan_inf(self, obj, recursion_depth):
+        if obj == float("inf"):
+            expected = "float('inf')"
+        elif obj == float("-inf"):
+            expected = "float('-inf')"
+        else:
+            expected = "float('nan')"
+
+        for _ in range(recursion_depth):
+            obj = {"x": obj}
+            expected = f"{{'x': {expected}}}"
+
+        assert util.make_repr("Example", obj) == (f"Example({expected})", False, True)
+
+
+@pytest.mark.serial
+def test_check_called_by():
+    outfile = io.StringIO()
+    with contextlib.redirect_stdout(outfile):
+        warn_msg = "Calling 'test_check_called_by.<locals>.callee()' directly is not recommended. Please use 'caller()' instead."
+
+        @util.check_called_by("caller")
+        def callee():
+            pass
+
+        def caller():
+            return callee()
+
+        # If we call via the caller, no message should be emitted
+        caller()
+        outfile.seek(0)
+        out = outfile.read()
+        assert warn_msg not in out
+
+        # If we call the callee directly, we should see a warning
+        callee()
+        outfile.seek(0)
+        out = outfile.read()
+        assert warn_msg in out
+
+
+class TestGetNumBytes:
+    def test_should_get_given_str(self) -> None:
+        """Test that _get_num_bytes returns the correct number of bytes when given `str`."""
+        # Precondition.
+        contents = "hello"
+
+        # Under test.
+        num_bytes = util_internal._get_num_bytes(contents)
+
+        # Postcondition.
+        assert num_bytes == len("hello")
+
+    def test_should_get_given_bytes(self) -> None:
+        """Test that _get_num_bytes returns the correct number of bytes when given `bytes`."""
+        # Precondition.
+        contents = bytes(b"hello")
+
+        # Under test.
+        num_bytes = util_internal._get_num_bytes(contents)
+
+        # Postcondition.
+        assert num_bytes == len("hello")
+
+    def test_should_get_given_IHostMemory(self) -> None:
+        """Test that _get_num_bytes returns the correct number of bytes when given `IHostMemory`."""
+        # Precondition.
+        contents = engine_from_network(network_from_onnx_bytes(ONNX_MODELS["identity"].loader)).serialize()
+
+        # Under test.
+        num_bytes = util_internal._get_num_bytes(contents)
+
+        # Postcondition.
+        assert num_bytes == len(memoryview(contents))
+
+    def test_should_raise_error_given_invalid_type(self) -> None:
+        """Test that _get_num_bytes raises an error when given an invalid type."""
+        # Precondition.
+        invalid_contents = 123
+
+        # Under test and postcondition.
+        with pytest.raises(
+            TypeError, match=f"`contents` is {invalid_contents}, which is not bytes-like. Cannot get number of bytes."
+        ):
+            util_internal._get_num_bytes(invalid_contents)
+
+
+class TestSaveFile:
+    def test_should_save_str_to_path(self) -> None:
+        """Test that `save_file` should save a string to a path."""
+        # Precondition.
+        contents = "hello"
+        with util.NamedTemporaryFile("w+") as f:
+            dest = f.name
+
+            # Under test.
+            util.save_file(contents, dest, mode="w+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert contents == written_contents
+
+    def test_should_save_str_to_file_like(self) -> None:
+        """Test that `save_file` should save a string to a file-like object."""
+        # Precondition.
+        contents = "hello"
+        with util.NamedTemporaryFile("w+") as f:
+            dest = f
+
+            # Under test.
+            util.save_file(contents, dest, mode="w+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert contents == written_contents
+
+    def test_should_save_bytes_to_path(self) -> None:
+        """Test that `save_file` should save bytes to a path."""
+        # Precondition.
+        contents = b"hello"
+        with util.NamedTemporaryFile("wb+") as f:
+            dest = f.name
+
+            # Under test.
+            util.save_file(contents, dest, mode="wb+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert contents == written_contents
+
+    def test_should_save_IHostMemory_to_path(self) -> None:
+        """Test that `save_file` should save an `IHostMemory` to a path."""
+        # Precondition.
+        contents = engine_from_network(network_from_onnx_bytes(ONNX_MODELS["identity"].loader)).serialize()
+        with util.NamedTemporaryFile("wb+") as f:
+            dest = f.name
+
+            # Under test.
+            util.save_file(contents, dest, mode="wb+")
+
+            # Postcondition.
+            f.seek(0)
+            written_contents = f.read()
+            assert bytes(contents) == written_contents