Python Target 参数化

Python Target 参数化#

概要#

对于任何支持的运行时环境，TVM 都应生成数值上正确的结果。因此，在编写验证数值输出的单元测试时，这些测试应在所有支持的运行时环境中运行。由于这是非常常见的用例，TVM 提供了辅助函数来参数化单元测试，使得它们能够在所有启用且设备兼容的目标上运行。

测试套件中的单独的 Python 函数可以扩展为多个参数化的单元测试，每个测试针对特定的目标设备。要执行测试，以下所有条件必须成立。

  • 这个测试存在于文件或目录中，该文件或目录已被传递到 pytest。

  • 应用到函数上的 pytest 标记，无论是显式地还是通过目标参数化的方式，都必须与传递给 pytest 的 -m 参数表达式兼容。

  • 对于使用 target 装置的参数化测试，目标必须出现在环境变量 TVM_TEST_TARGETS 中。

  • 对于使用 target fixture的参数化测试，必须在 config.cmake 中的构建配置中启用相应的运行时环境。

单元测试文件内容#

在多个目标上运行测试的推荐方法是通过参数化测试。对于固定的目标列表，可以通过使用 @tvm.testing.parametrize_targets('target_1', 'target_2', ...) 装饰器显式地实现，并在函数中接受 target 或 dev 作为参数。该函数将为列出的每个目标运行一次，并且每个目标的成功或失败将单独报告。如果某个目标无法运行，例如在 config.cmake 中被禁用，或者没有合适的硬件存在，则该目标将被报告为跳过。

# Explicit listing of targets to use.
@tvm.testing.parametrize_target('llvm', 'cuda')
def test_function(target, dev):
    # Test code goes here

对于那些应当能在所有目标上正确运行的测试，可以省略装饰器。任何接受 target 或 dev 参数的测试都会自动针对 TVM_TEST_TARGETS 中指定的所有目标进行参数化。这种参数化方式为每个目标提供相同的通过/失败/跳过报告，同时使得测试套件能够轻松扩展以覆盖更多目标。

# Implicitly parametrized to run on all targets
# in environment variable TVM_TEST_TARGETS
def test_function(target, dev):
    # Test code goes here

@tvm.testing.parametrize_targets 也可以作为简单的装饰器使用，以明确提示测试已被参数化，但这并不会产生额外的效果。

# Explicitly parametrized to run on all targets
# in environment variable TVM_TEST_TARGETS
@tvm.testing.parametrize_targets
def test_function(target, dev):
    # Test code goes here

可以使用 @tvm.testing.exclude_targets 或 @tvm.testing.known_failing_targets 装饰器来排除特定目标或将其标记为预期失败。有关其预期用例的更多信息，请参阅它们的文档字符串。

在某些情况下，可能需要对多个参数进行参数化。例如，可能存在需要测试的特定目标实现，其中某些目标拥有不止一种实现方式。这可以通过显式地对参数元组进行参数化来完成，如下所示。在这些情况下，只有明确列出的目标会运行，但它们仍然会应用适当的 @tvm.testing.requires_RUNTIME 标记。

@pytest.mark.parametrize('target,impl', [
     ('llvm', cpu_implementation),
     ('cuda', gpu_implementation_small_batch),
     ('cuda', gpu_implementation_large_batch),
 ])
 def test_function(target, dev, impl):
     # Test code goes here

参数化功能是基于 pytest 标记实现的。每个测试函数都可以用 pytest标记 来装饰，以包含元数据。最常用的标记如下。

  • @pytest.mark.gpu - 将函数标记为使用 GPU 功能。这个标记本身不会产生任何效果，但可以与命令行参数 -m gpu 或 -m 'not gpu' 结合使用，以限制 pytest 执行的测试范围。这个标记不应单独使用，而是作为单元测试中其他标记的一部分。

  • @tvm.testing.uses_gpu - 应用了 @pytest.mark.gpu。这个装饰器应用于标记可能使用 GPU 的单元测试，如果存在 GPU 的话。这个装饰器仅对于显式循环遍历 tvm.testing.enabled_targets() 的测试是必要的，但这已不再是编写单元测试的首选风格（见下文）。当使用 tvm.testing.parametrize_targets() 时，对于 GPU 目标，这个装饰器是隐含的，不需要显式应用。

  • @tvm.testing.requires_gpu - 应用了 @tvm.testing.uses_gpu，并且额外标记了如果不存在 GPU，则测试应完全跳过（@pytest.mark.skipif）。

  • @tvm.testing.requires_RUNTIME - 多个装饰器（例如 @tvm.testing.requires_cuda），每个装饰器在指定的运行时无法使用时跳过测试。如果运行时在 config.cmake 中被禁用，或者没有兼容的设备存在，则运行时无法使用。对于使用 GPU 的运行时，这包括 @tvm.testing.requires_gpu。

当使用参数化目标时，每次测试运行都会根据所使用的目标自动应用相应的 @tvm.testing.requires_RUNTIME 装饰器。因此，如果某个目标在 config.cmake 中被禁用，或者没有适当的硬件来运行，它将被明确列为跳过状态。

还存在 tvm.testing.enabled_targets() 函数，它根据环境变量 TVM_TEST_TARGETS、构建配置和当前机器的物理硬件，返回所有已启用且可在当前机器上运行的目标。大多数现有测试显式地循环遍历从 enabled_targets() 返回的目标，但不建议在新测试中使用这种方式。这种风格的pytest输出会静默跳过在``config.cmake``中被禁用的运行时，或者没有设备可以运行的运行时。此外，测试会在第一个失败的目标处停止，这导致无法明确错误是发生在特定目标上，还是发生在所有目标上。

# Old style, do not use.
def test_function():
    for target,dev in tvm.testing.enabled_targets():
        # Test code goes here

本地运行#

要在本地运行 Python 单元测试，请在 ${TVM_HOME} 目录下使用命令 pytest。

  • 环境变量

    • TVM_TEST_TARGETS 应设置为以分号分隔的目标列表，用于指定要运行的目标。如果未设置，则默认为 tvm.testing.DEFAULT_TEST_TARGETS 中定义的目标。

      注意：如果 TVM_TEST_TARGETS 中不包含任何已启用且具有可访问设备的类型的目标，则测试将回退到仅在 llvm 目标上运行。

    • TVM_LIBRARY_PATH 应设置为指向 libtvm.so 库的路径。例如，这可以用于运行使用调试构建的测试。如果未设置，则会相对于TVM源代码目录搜索 libtvm.so。

  • 命令行参数

    • 传递文件夹或文件的路径将仅运行该文件夹或文件中的单元测试。例如，在没有安装特定前端的情况下，这可以避免运行位于 tests/python/frontend 中的测试。

    • -m 参数仅运行带有特定 pytest 标记的单元测试。最常见的用法是使用 -m gpu 来仅运行标记为 @pytest.mark.gpu 的测试，这些测试会使用 GPU 运行。它也可以用于仅运行不使用 GPU 的测试，通过传递 -m 'not gpu' 来实现。

      注意：此筛选操作是在基于 TVM_TEST_TARGETS 环境变量选择目标之后进行的。即使指定了 -m gpu，如果 TVM_TEST_TARGETS 中不包含 GPU 目标，则不会运行任何 GPU 测试。

在本地 Docker 容器中运行#

docker/bash.sh 脚本可用于在与 CI 使用的相同的 Docker 镜像中运行单元测试。第一个参数应指定要运行的 Docker 镜像（例如 docker/bash.sh ci_gpu）。允许的镜像名称定义在 TVM 源代码目录下的 Jenkinsfile 文件的顶部，并映射到 tlcpack 上的镜像。

如果没有提供额外的参数，Docker 镜像将会加载交互式的 bash 会话。如果传递了脚本作为可选参数（例如 docker/bash.sh ci_gpu tests/scripts/task_python_unittest.sh），那么该脚本将会在 Docker 镜像内部执行。

注意：Docker 镜像包含了所有的系统依赖项，但并未包含针对这些系统的 build/config.cmake 配置文件。TVM 源代码目录被用作 Docker 镜像的主目录，因此默认情况下会使用与本地配置相同的构建目录。一种解决方案是维护单独的 build_local 和 build_docker 目录，并在进入或退出 Docker 时，将 build 符号链接到相应的文件夹。

在 CI 中运行#

CI 中的所有流程都始于 Jenkinsfile 中的任务定义。这包括定义所使用的 Docker 镜像、编译时的配置以及各个阶段包含哪些测试。

  • Docker images

    Jenkinsfile 中的每个任务（例如 'BUILD: CPU'）都会调用 docker/bash.sh。调用 docker/bash.sh 后跟随的参数定义了 CI 中使用的 Docker 镜像，这与本地操作一致。

  • 编译时配置

    Docker 镜像中并未内置 config.cmake 文件，因此这是每个 BUILD 任务的第一步。这一步通过使用 tests/scripts/task_config_build_*.sh 脚本来完成。具体使用哪个脚本取决于正在测试的构建类型，并在 Jenkinsfile 中指定。

    每个 BUILD 任务最终都会打包为库，以供后续测试使用。

  • 运行哪些测试

    Jenkinsfile 中的 Unit Test 和 Integration Test 阶段决定了如何调用 pytest。每个任务首先解压在 BUILD 阶段预先编译好的库，然后运行测试脚本（例如 tests/script/task_python_unittest.sh）。这些脚本设置了传递给 pytest 的文件/文件夹和命令行选项。

    其中一些脚本包含了 -m gpu 选项，该选项限制了测试仅运行包含 @pytest.mark.gpu 标记的测试。