chore: import upstream snapshot with attribution

tools/Polygraphy/tests/util/test_serde.py

@@ -0,0 +1,218 @@
+#
+# 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 numpy as np
+import pytest
+import tensorrt as trt
+import torch
+
+from polygraphy import constants, util
+from polygraphy.backend.trt import Algorithm, TacticReplayData, TensorInfo
+from polygraphy.comparator import IterationResult, RunResults
+from polygraphy.exception import PolygraphyException
+from polygraphy.json import Decoder, Encoder, from_json, load_json, to_json
+
+
+class Dummy:
+    def __init__(self, x):
+        self.x = x
+
+
+@Encoder.register(Dummy)
+def encode_dummy(dummy):
+    return {"x": dummy.x}
+
+
+@Decoder.register(Dummy)
+def decode_dummy(dct):
+    assert len(dct) == 1  # Custom type markers should be removed at this point
+    return Dummy(x=dct["x"])
+
+
+class NoDecoder:
+    def __init__(self, x):
+        self.x = x
+
+
+@Encoder.register(NoDecoder)
+def encode_nodecoder(no_decoder):
+    return {"x": no_decoder.x}
+
+
+class TestEncoder:
+    def test_registered(self):
+        d = Dummy(x=-1)
+        d_json = to_json(d)
+        assert encode_dummy(d) == {"x": d.x, constants.TYPE_MARKER: "Dummy"}
+        expected = f'{{\n    "x": {d.x},\n    "{constants.TYPE_MARKER}": "Dummy"\n}}'
+        assert d_json == expected
+
+
+class TestDecoder:
+    def test_object_pairs_hook(self):
+        d = Dummy(x=-1)
+        d_json = to_json(d)
+
+        new_d = from_json(d_json)
+        assert new_d.x == d.x
+
+    def test_error_on_no_decoder(self):
+        d = NoDecoder(x=1)
+        d_json = to_json(d)
+
+        with pytest.raises(
+            PolygraphyException,
+            match="Could not decode serialized type: NoDecoder. This could be because a required module is missing.",
+        ):
+            from_json(d_json)
+
+    def test_names_correct(self):
+        # Trigger `try_register_common_json`
+        d = Dummy(x=-1)
+        to_json(d)
+
+        # If the name of a class changes, then we need to specify an `alias` when registering
+        # to retain backwards compatibility.
+        assert set(Decoder.polygraphy_registered.keys()) == {
+            "__polygraphy_encoded_Algorithm",
+            "__polygraphy_encoded_AttentionLayerHint",
+            "__polygraphy_encoded_Dummy",
+            "__polygraphy_encoded_FormattedArray",
+            "__polygraphy_encoded_IterationContext",
+            "__polygraphy_encoded_IterationResult",
+            "__polygraphy_encoded_LazyArray",
+            "__polygraphy_encoded_ndarray",
+            "__polygraphy_encoded_RunResults",
+            "__polygraphy_encoded_ShardHints",
+            "__polygraphy_encoded_ShardTensor",
+            "__polygraphy_encoded_TacticReplayData",
+            "__polygraphy_encoded_Tensor",
+            "__polygraphy_encoded_TensorInfo",
+            "Algorithm",
+            "AttentionLayerHint",
+            "Dummy",
+            "FormattedArray",
+            "IterationContext",
+            "IterationResult",
+            "LazyArray",
+            "LazyNumpyArray",
+            "ndarray",
+            "RunResults",
+            "ShardHints",
+            "ShardTensor",
+            "TacticReplayData",
+            "Tensor",
+            "TensorInfo",
+        }
+
+
+def make_algo():
+    return Algorithm(
+        implementation=4,
+        tactic=5,
+        # Should work even if strides are not set
+        inputs=[
+            TensorInfo(trt.float32, (1, 2), -1, 1),
+            TensorInfo(trt.float32, (1, 2), -1, 1),
+        ],
+        outputs=[TensorInfo(trt.float32, (2, 3), -1, 1)],
+    )
+
+
+def make_iter_result():
+    return IterationResult(
+        runtime=4.5,
+        runner_name="test",
+        outputs={
+            "out0": np.random.random_sample((1, 2, 1)),
+            "out1": np.ones((1, 2), dtype=np.float32),
+        },
+    )
+
+
+JSONABLE_CASES = [
+    RunResults([("runner0", [make_iter_result()]), ("runner0", [make_iter_result()])]),
+    TacticReplayData().add("hi", algorithm=make_algo()),
+]
+
+
+class TestImplementations:
+    @pytest.mark.parametrize(
+        "obj",
+        [
+            TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+            Algorithm(
+                implementation=4,
+                tactic=5,
+                inputs=[TensorInfo(trt.float32, (1, 2, 3), -1, 1)],
+                outputs=[TensorInfo(trt.float32, (1, 2, 3), -1, 1)],
+            ),
+            Algorithm(
+                implementation=4,
+                tactic=5,
+                inputs=[
+                    TensorInfo(trt.float32, (1, 2, 3), -1, 1),
+                    TensorInfo(trt.int8, (1, 2, 3), -1, 1),
+                ],
+                outputs=[TensorInfo(trt.float16, (1, 2, 3), -1, 1)],
+            ),
+            np.ones((3, 4, 5), dtype=np.int64),
+            np.ones(5, dtype=np.int64),
+            np.zeros((4, 5), dtype=np.float32),
+            np.random.random_sample((3, 5)),
+            torch.ones((3, 4, 5), dtype=torch.int64),
+            make_iter_result(),
+            RunResults(
+                [("runner0", [make_iter_result()]), ("runner0", [make_iter_result()])]
+            ),
+        ],
+        ids=lambda x: type(x),
+    )
+    def test_serde(self, obj):
+        encoded = to_json(obj)
+        decoded = from_json(encoded)
+        if isinstance(obj, np.ndarray):
+            assert np.array_equal(decoded, obj)
+        elif isinstance(obj, torch.Tensor):
+            assert torch.equal(decoded, obj)
+        else:
+            assert decoded == obj
+
+    @pytest.mark.parametrize("obj", JSONABLE_CASES)
+    def test_to_from_json(self, obj):
+        encoded = obj.to_json()
+        decoded = type(obj).from_json(encoded)
+        assert decoded == obj
+
+    @pytest.mark.parametrize("obj", JSONABLE_CASES)
+    def test_save_load(self, obj):
+        with util.NamedTemporaryFile("w+") as f:
+            obj.save(f)
+            decoded = type(obj).load(f)
+            assert decoded == obj
+
+    def test_cannot_save_load_to_different_types(self):
+        run_result = JSONABLE_CASES[0]
+        encoded = run_result.to_json()
+
+        with pytest.raises(PolygraphyException, match="JSON cannot be decoded into"):
+            TacticReplayData.from_json(encoded)
+
+
+def test_load_json_errors_if_file_nonexistent():
+    with pytest.raises(FileNotFoundError, match="No such file"):
+        load_json("polygraphy-nonexistent-path")