chore: import upstream snapshot with attribution

tools/Polygraphy/tests/mod/conftest.py

@@ -0,0 +1,29 @@
+#
+# 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 pytest
+from tests.helper import ROOT_DIR
+
+
+@pytest.fixture()
+def poly_venv(virtualenv):
+    virtualenv.env["PYTHONPATH"] = ROOT_DIR
+    virtualenv.env["LD_LIBRARY_PATH"] = ""
+
+    # Newer versions of setuptools break pytest-virtualenv
+    virtualenv.run([virtualenv.python, "-m", "pip", "install", "setuptools==59.6.0"])
+
+    return virtualenv

tools/Polygraphy/tests/mod/test_dependencies.py

@@ -0,0 +1,273 @@
+#
+# 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)

tools/Polygraphy/tests/mod/test_exporter.py

@@ -0,0 +1,228 @@
+#
+# 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 pytest
+from polygraphy import mod
+from polygraphy.backend.base import BaseLoader
+
+# For test_funcify_with_collision
+functor2 = None
+
+
+class TestExporter:
+    def test_func(self):
+        @mod.export()
+        def test_func0():
+            pass
+
+        assert "test_func0" in __all__
+
+    def test_class(self):
+        @mod.export()
+        class TestClass0:
+            pass
+
+        assert "TestClass0" in __all__
+
+    def test_funcify_func_fails(self):
+        with pytest.raises(AssertionError, match="must be a loader"):
+
+            @mod.export(funcify=True)
+            def test_func1():
+                pass
+
+    def test_funcify_non_base_loader_class(self):
+        with pytest.raises(AssertionError, match="must derive from BaseLoader"):
+
+            @mod.export(funcify=True)
+            class NonFunctor0:
+                def __init__(self, x):
+                    self.x = x
+
+    def test_funcify_duplicate_parameters_in_call_init(self):
+        with pytest.raises(
+            AssertionError, match="call_impl and __init__ have the same argument names"
+        ):
+
+            @mod.export(funcify=True)
+            class DupArgs(BaseLoader):
+                def __init__(self, x):
+                    self.x = x
+
+                def call_impl(self, x):
+                    self.x = x
+
+    def test_funcify_takes_docstring(self):
+        @mod.export(funcify=True)
+        class DocstringFunctor(BaseLoader):
+            """This is a docstring"""
+
+            def __init__(self):
+                pass
+
+            def call_impl(self):
+                pass
+
+        assert "DocstringFunctor" in __all__
+        assert "docstring_functor" in __all__
+
+        assert (
+            docstring_functor.__doc__
+            == "Immediately evaluated functional variant of :class:`DocstringFunctor` .\n"
+        )
+
+    def test_funcify_functor_no_call_args(self):
+        @mod.export(funcify=True)
+        class Functor0(BaseLoader):
+            def __init__(self, x):
+                self.x = x
+
+            def call_impl(self):
+                return self.x
+
+        assert "Functor0" in __all__
+        assert "functor0" in __all__
+        assert functor0(0) == 0
+
+    def test_funcify_functor_with_call_args(self):
+        @mod.export(funcify=True)
+        class Functor1(BaseLoader):
+            def __init__(self, x):
+                self.x = x
+
+            def call_impl(self, y, z):
+                return self.x, y, z
+
+        assert "Functor1" in __all__
+        assert "functor1" in __all__
+
+        # __init__ arguments always precede __call__ arguments
+        x, y, z = functor1(0, 1, -1)
+        assert (x, y, z) == (0, 1, -1)
+
+        # Keyword arguments should behave as expected
+        x, y, z = functor1(y=1, x=0, z=-1)
+        assert (x, y, z) == (0, 1, -1)
+
+    def test_funcify_functor_with_call_args_defaults(self):
+        @mod.export(funcify=True)
+        class FunctorWithCallArgs(BaseLoader):
+            def __init__(self, x=0):
+                self.x = x
+
+            def call_impl(self, y, z=-1):
+                return self.x, y, z
+
+        assert "FunctorWithCallArgs" in __all__
+        assert "functor_with_call_args" in __all__
+
+        # __init__ arguments always precede __call__ arguments
+        x, y, z = functor_with_call_args(y=1)
+        assert (x, y, z) == (0, 1, -1)
+
+        # Keyword arguments should behave as expected
+        x, y, z = functor_with_call_args(y=1)
+        assert (x, y, z) == (0, 1, -1)
+
+    def test_funcify_with_collision(self):
+        with pytest.raises(AssertionError, match="symbol is already defined"):
+
+            @mod.export(funcify=True)
+            class Functor2(BaseLoader):
+                def __init__(self, x):
+                    self.x = x
+
+                def call_impl(self, y, z):
+                    return self.x, y, z
+
+    def test_funcify_functor_with_dynamic_call_args_kwargs(self):
+        @mod.export(funcify=True)
+        class Functor3(BaseLoader):
+            def __init__(self, f):
+                self.f = f
+
+            def call_impl(self, *args, **kwargs):
+                return self.f(*args, **kwargs)
+
+        assert "Functor3" in __all__
+        assert "functor3" in __all__
+
+        # We should be able to pass arbitrary arguments to call now.
+        # __init__ arguments are always first.
+        def func(arg0, arg1, arg2):
+            return arg0 + arg1 + arg2
+
+        assert functor3(func, 1, 2, arg2=4) == 7
+
+    def test_funcify_with_inherited_init(self):
+        class BaseFunctor4(BaseLoader):
+            def __init__(self, x):
+                self.x = x
+
+        @mod.export(funcify=True)
+        class Functor4(BaseFunctor4):
+            def call_impl(self):
+                return self.x
+
+        assert "Functor4" in __all__
+        assert "functor4" in __all__
+
+        assert functor4(-1) == -1
+
+    def test_funcify_functor_with_default_vals(self):
+        @mod.export(funcify=True)
+        class FunctorWithDefaults(BaseLoader):
+            def __init__(self, w, x=1):
+                self.w = w
+                self.x = x
+
+            def call_impl(self, y, z=3):
+                return self.w, self.x, y, z
+
+        assert "FunctorWithDefaults" in __all__
+        assert "functor_with_defaults" in __all__
+
+        # Since x and z have default values, the arguments will be interlaced into:
+        # w, y, x, z
+        # __init__ parameters take precedence, and call_impl parameters follow.
+        w, x, y, z = functor_with_defaults(-1, -2)  # Set just w, y
+        assert (w, x, y, z) == (-1, 1, -2, 3)
+
+        w, x, y, z = functor_with_defaults(0, 1, 2, 3)  # Set all
+        assert (w, x, y, z) == (0, 2, 1, 3)
+
+    def test_funcify_functor_with_type_annotations(self):
+        @mod.export(funcify=True)
+        class FunctorWithTypeAnnotations(BaseLoader):
+            def __init__(self, w: int, x: int = 1):
+                self.w = w
+                self.x = x
+
+            def call_impl(self, y: int, z: int = 3):
+                return self.w, self.x, y, z
+
+        assert "FunctorWithTypeAnnotations" in __all__
+        assert "functor_with_type_annotations" in __all__
+
+        # Since x and z have default values, the arguments will be interlaced into:
+        # w, y, x, z
+        # __init__ parameters take precedence, and call_impl parameters follow.
+        w, x, y, z = functor_with_type_annotations(-1, -2)  # Set just w, y
+        assert (w, x, y, z) == (-1, 1, -2, 3)
+
+        w, x, y, z = functor_with_type_annotations(0, 1, 2, 3)  # Set all
+        assert (w, x, y, z) == (0, 2, 1, 3)

tools/Polygraphy/tests/mod/test_importer.py

@@ -0,0 +1,185 @@
+#
+# 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 copy
+import os
+import sys
+from textwrap import dedent
+
+import pytest
+import tempfile
+import tensorrt as trt
+from polygraphy import mod, util
+from polygraphy.exception import PolygraphyException
+from polygraphy.mod.importer import _version_ok
+
+common_backend = mod.lazy_import("polygraphy.backend.common")
+
+class TestImporter:
+    def test_import_from_script(self):
+        script = dedent(
+            """
+            from polygraphy.backend.trt import CreateNetwork
+            from polygraphy import func
+            import tensorrt as trt
+
+            @func.extend(CreateNetwork())
+            def load_network(builder, network):
+                inp = network.add_input("input", dtype=trt.float32, shape=(1, 1))
+                out = network.add_identity(inp).get_output(0)
+                network.mark_output(out)
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            orig_sys_path = copy.deepcopy(sys.path)
+            load_network = mod.import_from_script(f.name, "load_network")
+            assert sys.path == orig_sys_path
+            builder, network = load_network()
+            with builder, network:
+                assert isinstance(builder, trt.Builder)
+                assert isinstance(network, trt.INetworkDefinition)
+                assert network.num_layers == 1
+                assert network.get_layer(0).type == trt.LayerType.IDENTITY
+            assert sys.path == orig_sys_path
+
+    def test_import_from_script_same_method_different_modules(self):
+        module1_script = dedent(
+            """
+            def print_message():
+                print(f"msg1::print_message")
+                return "msg1"
+            """
+        )
+
+        module2_script = dedent(
+            """
+            def print_message():
+                print(f"msg2::print_message")
+                return "msg2"
+            """
+        )
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            os.mkdir(os.path.join(tempdir, "msg1"))
+            with open(os.path.join(tempdir, "msg1", "msg.py"), "w+") as msg1_msg:
+                msg1_msg.write(module1_script)
+                msg1_msg.flush()
+                os.fsync(msg1_msg.fileno())
+
+                os.mkdir(os.path.join(tempdir, "msg2"))
+                with open(os.path.join(tempdir, "msg2", "msg.py"), "w+") as msg2_msg:
+                    msg2_msg.write(module2_script)
+                    msg2_msg.flush()
+                    os.fsync(msg2_msg.fileno())
+                    
+                    for msg_module in ['msg1', 'msg2']:
+                        msg_loc = os.path.join(tempdir,msg_module,'msg.py')
+                        msg = common_backend.invoke_from_script(msg_loc, "print_message")
+                        assert msg==msg_module
+
+    def test_import_non_existent(self):
+        script = dedent(
+            """
+            def example():
+                pass
+            """
+        )
+
+        with util.NamedTemporaryFile("w+", suffix=".py") as f:
+            f.write(script)
+            f.flush()
+            os.fsync(f.fileno())
+
+            orig_sys_path = copy.deepcopy(sys.path)
+            example = mod.import_from_script(f.name, "example")
+            assert sys.path == orig_sys_path
+
+            assert example is not None
+            example()
+
+            with pytest.raises(
+                PolygraphyException, match="Could not import symbol: non_existent from"
+            ):
+                mod.import_from_script(f.name, "non_existent")
+            assert sys.path == orig_sys_path
+
+    @pytest.mark.parametrize(
+        "ver, pref, expected",
+        [
+            ("0.0.0", "==0.0.0", True),
+            ("0.0.0", "== 0.0.1", False),
+            ("0.0.0", ">= 0.0.0", True),
+            ("0.0.0", ">=0.0.1", False),
+            ("0.0.0", "<= 0.0.0", True),
+            ("0.0.2", "<=0.0.1", False),
+            ("0.0.1", "> 0.0.0", True),
+            ("0.0.1", ">0.0.1", False),
+            ("0.0.0", "< 0.0.1", True),
+            ("0.0.0", "< 0.0.0", False),
+            ("0.2.0", mod.LATEST_VERSION, False),
+        ],
+    )
+    def test_version_ok(self, ver, pref, expected):
+        assert _version_ok(ver, pref) == expected
+
+    def test_is_installed_works_when_package_name_differs_from_module_name(
+        self, poly_venv
+    ):
+        assert "onnxruntime" not in poly_venv.installed_packages()
+        assert "onnxruntime-gpu" not in poly_venv.installed_packages()
+
+        poly_venv.run(
+            [poly_venv.python, "-m", "pip", "install", "onnxruntime-gpu", "--no-deps"]
+        )
+
+        # The `onnxruntime-gpu` package provides the `onnxruntime` module.
+        # `is_installed()` should be able to understand that.
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                "from polygraphy import mod; onnxrt = mod.lazy_import('onnxruntime<0'); assert onnxrt.is_installed()",
+            ]
+        )
+
+    @pytest.mark.parametrize(
+        "mod_check", ["mod.has_mod('colored')", "colored.is_installed()"]
+    )
+    def test_has_mod(self, poly_venv, mod_check):
+        assert "colored" not in poly_venv.installed_packages()
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                f"from polygraphy import mod; colored = mod.lazy_import('colored'); assert not {mod_check}",
+            ]
+        )
+
+        poly_venv.run([poly_venv.python, "-m", "pip", "install", "colored==1.4.0"])
+        # Make sure `has_mod` doesn't actually import the package.
+        poly_venv.run(
+            [
+                poly_venv.python,
+                "-c",
+                "from polygraphy import mod; import sys; assert mod.has_mod('colored'); assert 'colored' not in sys.modules; import colored; assert 'colored' in sys.modules",
+            ]
+        )

tools/Polygraphy/tests/mod/test_util.py

@@ -0,0 +1,39 @@
+#
+# 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 pytest
+from polygraphy import mod
+
+
+@pytest.mark.parametrize(
+    "ver0, ver1, expected",
+    [
+        ("1.0.0", "2.0.0", False),
+        ("1.0.0", "0.9.0", True),
+        ("0.0b1", "0.1b1", False),
+        ("0.1b1", "0.1b0", True),
+        ("0.12b1", "0.1b0", True),
+        ("0.1b0", "0.1a0", True),
+        ("0.1rc0", "0.1b0", True),
+        ("0.post1", "0.post0", True),
+        ("0.post1", "0.post2", False),
+        ("1.13.1+cu117", "1.13.0", True),
+        ("1.13.1+cu117", "1.13.2", False),
+    ],
+)
+def test_version(ver0, ver1, expected):
+    assert (mod.version(ver0) > mod.version(ver1)) == expected