wehub-resource-sync
274 lines
9.1 KiB
Python
274 lines
9.1 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 glob
|
|
import os
|
|
import subprocess as sp
|
|
|
|
import sys
|
|
import pytest
|
|
from polygraphy import mod, util
|
|
from polygraphy.mod.importer import _version_ok
|
|
|
|
from tests.helper import ALL_TOOLS, POLYGRAPHY_CMD, ROOT_DIR
|
|
from tests.models.meta import ONNX_MODELS
|
|
|
|
"""
|
|
The tests here ensure that no additional dependencies are introduced into
|
|
the various modules under Polygraphy.
|
|
"""
|
|
|
|
|
|
def is_submodule(path):
|
|
file_mod = (
|
|
os.path.isfile(path)
|
|
and path.endswith(".py")
|
|
and os.path.basename(path) != "__init__.py"
|
|
)
|
|
dir_mod = os.path.isdir(path) and os.path.isfile(os.path.join(path, "__init__.py"))
|
|
return file_mod or dir_mod
|
|
|
|
|
|
MODULE_PATH = os.path.join(ROOT_DIR, "polygraphy")
|
|
SUBMODULE_PATHS = [
|
|
os.path.relpath(os.path.splitext(path)[0], ROOT_DIR)
|
|
for path in glob.iglob(os.path.join(MODULE_PATH, "**"), recursive=True)
|
|
if is_submodule(path)
|
|
]
|
|
|
|
|
|
class TestPublicImports:
|
|
def test_no_extra_submodule_dependencies_required(self, poly_venv):
|
|
# Submodules should not require any extra dependencies to import.
|
|
for submodule_path in SUBMODULE_PATHS:
|
|
submodule_name = ".".join(submodule_path.split(os.path.sep))
|
|
cmd = [poly_venv.python, "-c", f"from {submodule_name} import *"]
|
|
print(" ".join(map(str, cmd)))
|
|
output = poly_venv.run(cmd, capture=True)
|
|
print(output)
|
|
|
|
def test_can_json_without_numpy(self, poly_venv):
|
|
cmd = [
|
|
poly_venv.python,
|
|
"-c",
|
|
"from polygraphy.json import to_json, from_json; x = to_json(1); x = from_json(x)",
|
|
]
|
|
print(" ".join(map(str, cmd)))
|
|
output = poly_venv.run(cmd, capture=True)
|
|
print(output)
|
|
|
|
|
|
class TestToolImports:
|
|
# We should be able to at least launch tools with no dependencies installed.
|
|
@pytest.mark.parametrize("tool, subtools", ALL_TOOLS.items())
|
|
def test_can_run_tool_without_deps(self, poly_venv, tool, subtools):
|
|
BASE_TOOL_CMD = [poly_venv.python, *POLYGRAPHY_CMD, tool, "-h"]
|
|
|
|
def check_tool(tool):
|
|
output = poly_venv.run(tool, capture=True)
|
|
assert "This tool could not be loaded due to an error:" not in output
|
|
assert "error:" not in output
|
|
assert "could not be loaded" not in output
|
|
|
|
check_tool(BASE_TOOL_CMD)
|
|
|
|
for subtool in subtools:
|
|
check_tool(BASE_TOOL_CMD + [subtool])
|
|
|
|
|
|
class TestAutoinstallDeps:
|
|
@pytest.mark.parametrize(
|
|
"cmd",
|
|
[
|
|
["run", ONNX_MODELS["identity"].path, "--onnxrt"],
|
|
["run", ONNX_MODELS["identity"].path, "--trt"],
|
|
[
|
|
"surgeon",
|
|
"sanitize",
|
|
"--fold-constants",
|
|
ONNX_MODELS["const_foldable"].path,
|
|
"-o",
|
|
util.NamedTemporaryFile().name,
|
|
],
|
|
],
|
|
)
|
|
@pytest.mark.slow
|
|
def test_can_automatically_install_deps(self, poly_venv, cmd):
|
|
poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
|
|
cmd = [poly_venv.python, *POLYGRAPHY_CMD] + cmd
|
|
print(f"Running: {' '.join(map(str, cmd))}")
|
|
output = poly_venv.run(cmd, capture=True)
|
|
print(output)
|
|
assert "is required, but not installed. Attempting to install now" in output
|
|
|
|
# Make sure that no PyTorch dependency is accidentally introduced
|
|
assert "torch" not in poly_venv.installed_packages()
|
|
|
|
@pytest.mark.parametrize(
|
|
"new_ver, expected",
|
|
[
|
|
("==1.4.2", "==1.4.2"),
|
|
(mod.LATEST_VERSION, ">=1.4.2"),
|
|
],
|
|
)
|
|
def test_can_automatically_upgrade_deps(self, poly_venv, new_ver, expected):
|
|
poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
|
|
|
|
def get_colored_version():
|
|
return poly_venv.installed_packages()["colored"].version
|
|
|
|
poly_venv.run([poly_venv.python, "-m", "pip", "install", "colored==1.4.0"])
|
|
assert get_colored_version() == "1.4.0"
|
|
|
|
# Insert our own preferred version to make sure it upgrades.
|
|
poly_venv.run(
|
|
[
|
|
poly_venv.python,
|
|
"-c",
|
|
f"from polygraphy import mod; colored = mod.lazy_import('colored{new_ver}'); print(colored.__version__)",
|
|
]
|
|
)
|
|
assert _version_ok(get_colored_version(), expected)
|
|
|
|
# Make sure the `requires` parameter of `lazy_import` functions as we expect.
|
|
@pytest.mark.parametrize("preinstall", [True, False])
|
|
@pytest.mark.parametrize(
|
|
"new_ver, expected",
|
|
[
|
|
("==1.4.2", "==1.4.2"),
|
|
],
|
|
)
|
|
def test_can_automatically_install_requirements(
|
|
self, poly_venv, new_ver, expected, preinstall
|
|
):
|
|
poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
|
|
|
|
def get_colored_version():
|
|
return poly_venv.installed_packages()["colored"].version
|
|
|
|
if preinstall:
|
|
poly_venv.run([poly_venv.python, "-m", "pip", "install", "colored==1.4.0"])
|
|
assert get_colored_version() == "1.4.0"
|
|
|
|
# Insert our own preferred version to make sure it upgrades.
|
|
poly_venv.run(
|
|
[
|
|
poly_venv.python,
|
|
"-c",
|
|
f"from polygraphy import mod; "
|
|
f"requests = mod.lazy_import('requests==2.25.1', requires=['colored{new_ver}']); "
|
|
f"requests.__version__; "
|
|
f"import colored; print(colored.__version__)",
|
|
]
|
|
)
|
|
assert _version_ok(get_colored_version(), expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"import_params",
|
|
[
|
|
"'colored'",
|
|
"'colored', pkg_name='colored'",
|
|
"'colored', install_flags=['--force-reinstall']",
|
|
],
|
|
)
|
|
def test_autoinstall(self, poly_venv, import_params):
|
|
poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
|
|
assert "colored" not in poly_venv.installed_packages()
|
|
|
|
poly_venv.run(
|
|
[
|
|
poly_venv.python,
|
|
"-c",
|
|
f"from polygraphy import mod; colored = mod.lazy_import({import_params}); mod.autoinstall(colored)",
|
|
]
|
|
)
|
|
|
|
assert "colored" in poly_venv.installed_packages()
|
|
|
|
@pytest.mark.parametrize(
|
|
"response, should_install",
|
|
[
|
|
(" ", True),
|
|
("yes", True),
|
|
("Y", True),
|
|
("n", False),
|
|
],
|
|
)
|
|
def test_ask_before_autoinstall(self, response, should_install, poly_venv):
|
|
poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
|
|
poly_venv.env["POLYGRAPHY_ASK_BEFORE_INSTALL"] = "1"
|
|
assert "colored" not in poly_venv.installed_packages()
|
|
|
|
process = sp.Popen(
|
|
[
|
|
poly_venv.python,
|
|
"-c",
|
|
"from polygraphy import mod; "
|
|
"colored = mod.lazy_import('colored'); "
|
|
"mod.autoinstall(colored)",
|
|
],
|
|
env=poly_venv.env,
|
|
stdin=sp.PIPE,
|
|
)
|
|
process.communicate(input=response.encode())
|
|
process.wait()
|
|
|
|
assert ("colored" in poly_venv.installed_packages()) == should_install
|
|
|
|
# We can import inner modules, and Polygraphy should still autoinstall the outermost one.
|
|
def test_can_install_for_nested_import(self, poly_venv):
|
|
poly_venv.env["POLYGRAPHY_AUTOINSTALL_DEPS"] = "1"
|
|
|
|
poly_venv.run(
|
|
[
|
|
poly_venv.python,
|
|
"-c",
|
|
"from polygraphy import mod; "
|
|
"shape_inference = mod.lazy_import('onnx.shape_inference'); "
|
|
"print(shape_inference.infer_shapes)",
|
|
]
|
|
)
|
|
|
|
assert "onnx" in poly_venv.installed_packages()
|
|
|
|
def test_all_lazy_imports(self):
|
|
# NOTE: If this test fails, it means a new lazy dependency has been
|
|
# introduced. Please ensure that AUTOINSTALL continues to work with the
|
|
# new dependency.
|
|
expected = [
|
|
"fcntl",
|
|
"matplotlib.pyplot",
|
|
"matplotlib",
|
|
"msvcrt",
|
|
"numpy",
|
|
"onnx_graphsurgeon",
|
|
"onnx.numpy_helper",
|
|
"onnx",
|
|
"onnxmltools",
|
|
"onnxruntime.tools.symbolic_shape_infer",
|
|
"onnxruntime",
|
|
"tensorflow",
|
|
"tensorrt",
|
|
"tf2onnx",
|
|
"torch",
|
|
"yaml",
|
|
]
|
|
if sys.version_info < (3, 8):
|
|
expected.append("importlib_metadata")
|
|
|
|
assert mod.importer._all_external_lazy_imports == set(expected)
|