从零编写 Taichi Python 测试：pytest、多后端适配与 `@ti.test` 装饰器完全指南

从零编写 Taichi Python 测试pytest、多后端适配与ti.test装饰器完全指南【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi本篇指南以 Taichi 官方贡献文档docs/lang/articles/contribution/write_test.md为主线系统讲解如何在 Taichi 项目中编写、组织与运行 Python 功能测试从最基础的test_函数、0-D Field 传参到基于ti.test的多后端CPU/CUDA/Metal/OpenGL/Vulkan测试、ti.approx容差比较、pytest.mark.parametrize参数化再到扩展能力筛选与ti.init配置注入。读完本文你可以独立为 Taichi 的新特性例如一个新的数学函数写出覆盖多后端的健壮测试并借助仓库自带的测试工具链验证其正确性。1. 背景Taichi 的 Python 测试基础设施Taichi 是一个在 Python 中编写可移植高性能 GPU 程序的框架其功能测试主要用 Python 编写底层由 pytest 驱动。从 tests/run_tests.py 的源码可以看出仓库的测试入口本质上就是对 pytest 的封装测试文件统一放在 tests/python 目录下文件命名约定为test_xxx.py运行测试时通过 tests/run_tests.py 传入文件名关键字如logarithm脚本会把它自动补全为test_logarithm.py并交给 pytest 执行运行 Python 测试前需先安装测试依赖见 requirements_test.txtpip install -r requirements_test.txt依赖中除 pytest 本身外还包括pytest-xdist并行测试、pytest-rerunfailures失败重跑、pytest-cov覆盖率、numpy、nbmake支持在测试中运行 Notebook等它们共同支撑起测试的并行执行、超时控制与覆盖统计能力。tests/run_tests.py提供的常用命令行选项包括选项作用files位置参数指定要运行的测试文件关键字如python tests/run_tests.py logarithm-v / --verbose详细输出-k / --keys按关键字筛选测试-a / --arch指定运行的后端arch如cpu,cuda-n / --exclusive与-a配合排除指定后端-t / --threads自定义并行线程数默认取TI_TEST_THREADS环境变量或 8 与 CPU 核数的较小值-T / --timeout单个 Python 测试的超时时间默认 600 秒-C / --coverage收集覆盖率-x / --fail-fast首个失败即退出--with-offline-cache以离线缓存模式运行测试此外脚本还会根据TI_WANTED_ARCHS环境变量决定目标后端集合未设置时默认在机器上所有受支持的后端上运行。2. 添加第一个测试用例2.1 函数命名与文件组织pytest 会收集以test_开头的函数因此测试函数名必须以test_开头。假设你刚为 Taichi 新增了一个工具函数ti.log10想确保它的行为正确。首先查看 tests/python 下是否已有适合放置该测试的文件如果没有就新建一个例如tests/python/test_logarithm.pyimport taichi as ti def test_log10(): pass2.2 通过 0-D Field 与 Taichi 作用域交互在 Taichi 中Python 作用域与 Taichi 作用域ti.kernel内部之间不能直接读写普通 Python 变量通常借助 0-D Field即形状为()的 Field用r[None]访问来完成数值传递。一个完整的冒烟测试如下import taichi as ti def test_log10(): ti.init(archti.cpu) r ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.log10(r[None]) r[None] 100 foo() assert r[None] 2这段代码做了四件事ti.init(archti.cpu)初始化 Taichi 运行时指定在 CPU 后端执行声明 0-D Fieldr类型为ti.f32定义 kernelfoo在 Taichi 作用域内对r[None]原地取对数在 Python 作用域写入100、调用 kernel、再断言结果为2。2.3 运行测试python tests/run_tests.py logarithm该命令会定位到tests/python/test_logarithm.py并执行其中所有test_开头的函数。如果想更细致地控制也可以直接用 pytest 运行python -m pytest tests/python/test_logarithm.py -v注意本仓库对测试文件的自动补全逻辑文件名前缀test_、后缀.py实现在 tests/run_tests.py 的_test_python函数中因此传入logarithm而不是完整文件名是允许的。3. 在多个后端上运行测试上面例子里的ti.init(archti.cpu)把测试固定在了 CPU 后端。要让一个测试在多个后端上自动运行应使用ti.test装饰器。3.1 指定后端列表import taichi as ti # 会在 CPU 和 CUDA 两个后端上分别执行 ti.test(arch[ti.cpu, ti.cuda]) def test_log10(): r ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.log10(r[None]) r[None] 100 foo() assert r[None] 2装饰后pytest 会将该用例按每个后端参数化逐后端独立初始化并执行。注意此时不应再在函数体内调用ti.init()初始化参数统一交给ti.test管理。3.2 覆盖全部可用后端不指定arch参数即可在所有可用后端上测试import taichi as ti # 会在当前环境所有可用后端上执行 ti.test() def test_log10(): r ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.log10(r[None]) r[None] 100 foo() assert r[None] 23.3ti.test的底层行为ti.test的实现位于 tests/test_utils.py 的test函数其核心逻辑如下通过expected_archs()读取TI_WANTED_ARCHS环境变量通常由 tests/run_tests.py 的-a/-n选项设置得到目标后端集合用pytest.mark.parametrize把测试按(arch, options)展开每个用例的 id 形如archcpu-0、archcuda-1便于在失败输出中定位具体后端通过_ti_core.is_extension_supported检查后端是否支持所要求的扩展若某后端不满足条件直接跳过skip不会报错同时会把函数标记为func.__ti_test__ True供测试框架识别。仓库中 tests/python/test_test.py 本身就是对这套测试工具的自测test_multiple_archs断言当前运行的 arch 属于给定列表test_exclude_cpu断言exclude后 CPU 不再执行test_init_args验证ti.test能把debug、advanced_optimization等参数正确传入配置。写测试时可以参考这些现成范式。4. 用ti.approx处理跨后端精度差异4.1 问题场景部分后端典型如 OpenGL的浮点运算精度有限ti.log10(100)可能返回2.000001或1.999999。若用assert r[None] 2这样的精确比较这类后端上的测试会不稳定地失败。4.2 引入容差比较ti.approx提供了带容差的比较2.001 ti.approx(2)在 OpenGL 后端上返回Trueimport taichi as ti # 会在当前环境所有可用后端上执行 ti.test() def test_log10(): r ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.log10(r[None]) r[None] 100 foo() assert r[None] ti.approx(2)4.3 为什么不直接用pytest.approx原文档特别指出pytest.approx在这里并不合适原因是它的容差是固定的不会随 Taichi 后端的不同而变化在 OpenGL 这类低精度后端上很可能失败。ti.approx则针对 Taichi 各后端做了适配。从 tests/test_utils.py 的实现可以看到其原理内部基于pytest.approx实现但会根据当前后端动态抬高相对容差下限OpenGL 为1e-3Metal 为1e-4其余后端为1e-6见get_rel_eps对布尔类型做了特殊处理2 ti.approx(True)成立因为布尔值比较时只关心真值语义而非数值相等还提供了allclose(x, y, **kwargs)便捷函数等价于x approx(y, **kwargs)。因此凡是涉及浮点运算、且可能跑在不同精度后端上的测试都建议用ti.approx而非。5. 用pytest.mark.parametrize参数化输入上一节的测试只验证了输入100这一种情况。要覆盖多组输入可使用pytest.mark.parametrize这是 pytest 原生的参数化机制与ti.test可以叠加使用。5.1 单参数多取值import taichi as ti import pytest import math pytest.mark.parametrize(x, [1, 10, 100]) ti.test() def test_log10(x): r ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.log10(r[None]) r[None] x foo() assert r[None] math.log10(x)这里用 Python 标准库math.log10(x)作为期望值避免了手写常量也顺便覆盖了与 numpy 等参照实现的一致性验证思路。5.2 多参数成组取值多个参数可以打包成元组列表一起参数化import taichi as ti import pytest import math pytest.mark.parametrize(x,y, [(1, 2), (1, 3), (2, 1)]) ti.test() def test_atan2(x, y): r ti.field(ti.f32, ()) s ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.atan2(r[None]) r[None] x s[None] y foo() assert r[None] math.atan2(x, y)注意示例中 kernel 内对r[None]取atan2并写回自身仅用于演示参数化写法真实测试中应按需把s[None]传入运算。5.3 多参数全组合使用两个独立的parametrizepytest 会生成所有参数的笛卡尔积import taichi as ti import pytest import math pytest.mark.parametrize(x, [1, 2]) pytest.mark.parametrize(y, [1, 2]) # 等价于 .parametrize(x,y, [(1, 1), (1, 2), (2, 1), (2, 2)]) ti.test() def test_atan2(x, y): r ti.field(ti.f32, ()) s ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.atan2(r[None]) r[None] x s[None] y foo() assert r[None] math.atan2(x, y)这样共生成2 × 2 4组输入组合再乘以后端参数化一个用例即可完成对多输入、多后端的全面覆盖。6. 在ti.test中指定ti.init配置ti.test除了接收arch、exclude、require三个专有参数外其余关键字参数都会被透传给ti.init()。ti.test(ti.cpu, debugTrue, log_levelti.TRACE) def test_debugging_utils(): # ... (某些必须在调试模式下进行的测试)等价于def test_debugging_utils(): ti.init(archti.cpu, debugTrue, log_levelti.TRACE) # ... (某些必须在调试模式下进行的测试)也就是说ti.test会在每个后端用例初始化时替你调用ti.init(arch..., **options)。这在测试依赖调试模式、指定日志级别、关闭高级优化advanced_optimizationFalse等场景下非常有用且能保证每个后端独立获得一致的配置。该透传机制在 tests/test_utils.py 的test实现中通过copy.deepcopy(options)与参数化选项合并完成。7. 排除某些后端7.1 按后端排除某些特性在特定后端上不受支持例如 OpenGL 不支持稀疏数据结构可以用exclude参数跳过# 在所有后端上运行但 OpenGL 除外 ti.test(exclude[ti.opengl]) def test_sparse_field(): # ... (依赖 OpenGL 不支持的稀疏特性的测试)exclude的实现在 tests/test_utils.py 的exclude_arch_platform函数中支持三种形式单个后端excludeti.opengl后端列表exclude[ti.opengl, ti.metal]后端与平台的二元组exclude[(ti.vulkan, Darwin)]可在特定操作系统上排除特定后端7.2 按扩展能力排除如果只是想排除不具备某个特性的后端使用extensions关键字更语义化——ti.test会检查_ti_core.is_extension_supported(arch, ext)不支持该扩展的后端自动被跳过# 在所有支持 sparse 扩展的后端上运行 ti.test(extensions[ti.extension.sparse]) def test_sparse_field(): # ... (依赖 sparse 特性的测试)8. 要求测试依赖特定扩展require与排除相反若一个测试必须依赖某个扩展才能正确执行应在ti.test中加require参数。只有支持该扩展的后端才会运行此用例ti.test(requireti.extension.sparse) def test_struct_for_pointer_block(): n 16 block_size 8 f ti.field(dtypeti.f32) block ti.root.pointer(ti.ijk, n // block_size) block.dense(ti.ijk, block_size).place(f) f[0, 2, 3] 1 ti.kernel def count() - int: tot 0 for I in ti.grouped(block): tot 1 return tot assert count() 1这个例子完整演示了稀疏数据结构 结构化 for 循环遍历 返回整数的组合测试构建 pointer 块 dense 子块的 SNode 树写入单个元素后用ti.grouped遍历统计非空块数量。require支持单个扩展或扩展列表如require[ti.extension.sparse, ti.extension.bls]。在后端不支持所需扩展时用例会被自动跳过而非报错。当前支持的扩展清单require与extensions可用的扩展名定义在 taichi/inc/extensions.inc.h当前包括名称扩展内容sparse稀疏数据结构quant_basic量化中的基础操作quant完整量化功能meshMeshTaichi网格处理data6464 位数据与运算Metal 尚不支持 64 位数据缓冲区adstack自动微分中保存可变局部变量历史bls块本地存储Block-local storageassertionTaichi kernel 内的运行时断言extfunc支持插入外部函数调用或后端源码提示文档中的扩展表格与源码一一对应若后续新增扩展请同步更新 taichi/inc/extensions.inc.h。9. 综合实战把所学组合起来一个高质量的新特性测试往往同时用到多后端、参数化、容差比较与扩展约束。以ti.log10为例最终形态可以是import math import pytest import taichi as ti pytest.mark.parametrize(x, [1, 10, 100, 1000]) ti.test(arch[ti.cpu, ti.cuda], require[]) def test_log10(x): r ti.field(ti.f32, ()) ti.kernel def foo(): r[None] ti.log10(r[None]) r[None] x foo() assert r[None] ti.approx(math.log10(x))要点回顾测试文件放在 tests/python函数以test_开头用ti.test替代手动ti.init实现多后端参数化与初始化配置注入浮点结果一律用ti.approx比较避免 OpenGL 等低精度后端的不稳定失败用pytest.mark.parametrize扩展输入覆盖范围用math/numpy作为参照实现涉及后端特性差异时用exclude、extensions或require精确控制运行范围通过python tests/run_tests.py 关键字运行必要时配合-a指定后端、-v详细输出、-x失败即停等选项。10. 进一步探索测试工具实现ti.test、ti.approx、get_rel_eps等定义在 tests/test_utils.py测试入口与命令行选项tests/run_tests.py测试依赖清单requirements_test.txt测试工具自身的自测用例tests/python/test_test.py扩展清单require/extensions取值来源taichi/inc/extensions.inc.h官方贡献指南docs/lang/articles/contribution/write_test.md。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考