wehub-resource-sync
164 lines
5.4 KiB
Python
164 lines
5.4 KiB
Python
#
|
|
# 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 numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from polygraphy import config, util
|
|
from polygraphy.comparator import IterationResult, RunResults
|
|
from polygraphy.comparator.struct import LazyArray
|
|
from polygraphy.exception import PolygraphyException
|
|
|
|
|
|
def make_outputs():
|
|
return {"dummy_out": np.zeros((4, 4))}
|
|
|
|
|
|
def make_iter_results(runner_name):
|
|
return [IterationResult(outputs=make_outputs(), runner_name=runner_name)] * 2
|
|
|
|
|
|
@pytest.fixture()
|
|
def run_results():
|
|
results = RunResults()
|
|
results.append(("runner0", make_iter_results("runner0")))
|
|
results.append(("runner1", make_iter_results("runner1")))
|
|
return results
|
|
|
|
|
|
class TestRunResults:
|
|
def test_items(self, run_results):
|
|
for name, iteration_results in run_results.items():
|
|
assert isinstance(name, str)
|
|
assert isinstance(iteration_results, list)
|
|
for iter_res in iteration_results:
|
|
assert isinstance(iter_res, IterationResult)
|
|
|
|
def test_keys(self, run_results):
|
|
assert list(run_results.keys()) == ["runner0", "runner1"]
|
|
|
|
def test_values(self, run_results):
|
|
for iteration_results in run_results.values():
|
|
for iter_res in iteration_results:
|
|
assert isinstance(iter_res, IterationResult)
|
|
|
|
def test_getitem(self, run_results):
|
|
assert isinstance(run_results["runner0"][0], IterationResult)
|
|
assert isinstance(run_results[0][1][0], IterationResult)
|
|
assert run_results[0][1] == run_results["runner0"]
|
|
assert run_results[1][1] == run_results["runner1"]
|
|
|
|
def test_getitem_out_of_bounds(self, run_results):
|
|
with pytest.raises(IndexError):
|
|
run_results[2]
|
|
|
|
with pytest.raises(PolygraphyException, match="does not exist in this"):
|
|
run_results["runner2"]
|
|
|
|
def test_setitem(self, run_results):
|
|
def check_results(results, is_none=False):
|
|
for iter_res in results["runner1"]:
|
|
if is_none:
|
|
assert not iter_res
|
|
assert iter_res.runner_name == "custom_runner"
|
|
else:
|
|
assert iter_res
|
|
assert iter_res.runner_name
|
|
|
|
check_results(run_results)
|
|
|
|
iter_results = [IterationResult(outputs=None, runner_name=None)]
|
|
run_results["runner1"] = iter_results
|
|
|
|
check_results(run_results, is_none=True)
|
|
|
|
def test_setitem_out_of_bounds(self, run_results):
|
|
iter_results = [IterationResult(outputs=None, runner_name="new")]
|
|
run_results["runner2"] = iter_results
|
|
|
|
assert len(run_results) == 3
|
|
assert run_results["runner2"][0].runner_name == "new"
|
|
|
|
def test_contains(self, run_results):
|
|
assert "runner0" in run_results
|
|
assert "runner1" in run_results
|
|
assert "runner3" not in run_results
|
|
|
|
def test_add_new(self):
|
|
results = RunResults()
|
|
results.add([make_outputs()], runner_name="custom")
|
|
|
|
iter_results = results["custom"]
|
|
assert len(iter_results) == 1
|
|
assert all(
|
|
isinstance(iter_result, IterationResult) for iter_result in iter_results
|
|
)
|
|
|
|
def test_add_new_default_name(self):
|
|
results = RunResults()
|
|
results.add([make_outputs()])
|
|
|
|
name = results[0][0]
|
|
iter_results = results[name]
|
|
assert len(iter_results) == 1
|
|
assert all(
|
|
isinstance(iter_result, IterationResult) for iter_result in iter_results
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("module", [torch, np])
|
|
class TestLazyArray:
|
|
@pytest.mark.parametrize("set_threshold", [True, False])
|
|
def test_unswapped_array(self, set_threshold, module):
|
|
with contextlib.ExitStack() as stack:
|
|
if set_threshold:
|
|
|
|
def reset_array_swap():
|
|
config.ARRAY_SWAP_THRESHOLD_MB = -1
|
|
|
|
stack.callback(reset_array_swap)
|
|
|
|
config.ARRAY_SWAP_THRESHOLD_MB = 8
|
|
|
|
small_shape = (7 * 1024 * 1024,)
|
|
small_array = module.ones(small_shape, dtype=module.uint8)
|
|
lazy = LazyArray(small_array)
|
|
assert util.array.equal(small_array, lazy.arr)
|
|
assert lazy.tmpfile is None
|
|
|
|
assert util.array.equal(small_array, lazy.load())
|
|
|
|
def test_swapped_array(self, module):
|
|
with contextlib.ExitStack() as stack:
|
|
|
|
def reset_array_swap():
|
|
config.ARRAY_SWAP_THRESHOLD_MB = -1
|
|
|
|
stack.callback(reset_array_swap)
|
|
|
|
config.ARRAY_SWAP_THRESHOLD_MB = 8
|
|
|
|
large_shape = (9 * 1024 * 1024,)
|
|
large_array = module.ones(large_shape, dtype=module.uint8)
|
|
lazy = LazyArray(large_array)
|
|
assert lazy.arr is None
|
|
assert lazy.tmpfile is not None
|
|
|
|
assert util.array.equal(large_array, lazy.load())
|