博客

Pyrefly 现已支持 PyTorch 类型检查

我们很高兴地宣布,PyTorch 现已采用 Pyrefly 来支持核心代码库以及多个 PyTorch 生态系统项目(包括 Helion、TorchTitan 和 Ignite)的类型检查。对于 PyTorch 这种规模的项目,利用类型系统和类型检查长期以来对于确保一致性以及防止动态代码中常见的疏漏错误至关重要。迁移到 Pyrefly 为这些开发工作流带来了急需的升级,提供了闪电般快速、符合标准的类型检查和现代化的 IDE 体验。借助 Pyrefly,我们的维护者和贡献者可以更早地发现错误,从本地环境和 CI 运行的一致性结果中受益,并利用先进的类型功能。在这篇博文中,我们将分享我们做出这一转变的原因,并重点介绍 PyTorch 在采用 Pyrefly 后所获得的改进。

为什么要切换到 Pyrefly?

为了支持 PyTorch 的未来发展,我们需要一个速度快、易于使用、在不同开发环境中保持一致且得到积极维护的类型检查器。这些因素最终促使我们决定转向 Pyrefly。

平衡速度与准确性

在最近一轮对 PyTorch 进行类型检查的基准测试中,使用 MyPy 需要 50.6 秒,而 Pyrefly (v44.1) 仅需 5.5 秒。与 PyTorch 现有的工具相比,这在保持强大类型安全性的同时,速度有了显著提升。我们寻求的替代方案不仅要提供快速的结果,还要能帮助贡献者及早发现错误并识别类型覆盖范围中的空白。Pyrefly 似乎为我们找到了合适的平衡点,它足够快,可以跟上我们的开发进度,同时又不会牺牲类型检查的质量。

话虽如此,我们认为这仅仅是个开始;Pyrefly 仍有提升速度的空间,随着工具的不断演进,我们期望能获得更大的性能收益。我们将密切关注 Pyrefly 的持续开发,并期待在未来将其性能增强功能整合进来。

简化的配置

此前,我们对 MyPy 的依赖要求贡献者管理多个配置文件,以控制代码库中的覆盖率和严格程度。这使得确定哪些文件被检查以及使用哪些具体规则变得困难。迁移到 Pyrefly 有助于解决这些挑战。在 Pyrefly 团队的直接支持下,PyTorch 现已过渡到使用统一的 Pyrefly 配置和必要的抑制规则,这使得我们的维护者能够更轻松地了解哪些文件正在进行类型检查及其方式。 

跨开发环境的一致性

此前,开发者经常在 IDE、本地 CLI 和 CI 环境之间遇到差异,因为每个阶段使用的类型检查引擎不同。例如,PyTorch CI 作业可能使用 MyPy,但在 IDE 中,开发者倾向于使用表现略有不同的其他类型检查器。或者,开发者在 CLI 运行中启用的 MyPy 严格模式与 CI 中使用的模式不同。这些不一致导致了不可预测的反馈循环,并造成了令人沮丧的体验:即代码在本地类型检查中通过,但在 CI 中却失败了。通过采用同时提供高质量 IDE 体验以及强大 CLI 和 CI 功能的 Pyrefly,PyTorch 开发者现在可以在所有开发环境中获得一致的结果。

之前 之后
CI MyPy(完整项目运行) Pyrefly
CLI MyPy(仅限选定文件) Pyrefly
IDE Pyright 或其他 Pyrefly

积极的维护与快速开发

迁移的另一个主要原因是 Pyrefly 处于积极维护状态且发展迅速,具有持续性能改进的巨大空间。我们非常欣赏他们对用户反馈的响应能力和快速的开发周期,包括每周一发布新的次要版本。错误被报告并在下一次发布中得到解决的情况并不少见,这确保了问题能得到及时处理,新功能也能迅速交付。最近的一篇 Pyrefly 博文中描述了一个案例:一个性能瓶颈被识别并迅速解决,使得 PyTorch 代码库在 IDE 中的响应速度提升了 18 倍。

在整个迁移过程中,以及在我们继续使用 Pyrefly 的过程中,我们的首要任务是避免类型安全或开发者体验方面的倒退。与 Pyrefly 团队保持定期的沟通对于快速解决边缘情况并为贡献者提供顺利的过渡至关重要。

对 PyTorch 贡献者的额外好处

自转向 Pyrefly 以来,PyTorch 的贡献者和维护者已经体验到了显著的改进。除了迁移的最初动机外,其他好处还包括:

改进的代码质量

Pyrefly 的部署已经促成了 PyTorch 代码库中大量错误的发现和解决。实现这一目标的一个因素是 Pyrefly 在 PyTorch 中以一致的模式运行。以如下代码为例:除非 MyPy 处于严格模式,否则它不会对未类型化函数的主体进行类型检查,这意味着像这样的错误可能会被忽视。而 Pyrefly 在整个代码库中以一种统一的模式运行,能够捕获这些类型的错误。

def foo():
    return 1 + "" # pyrefly error

无缝的 IDE 体验

Pyrefly 与 许多主流 IDE 原生集成,将实时类型反馈、悬停文档和即时诊断直接带入编辑器,且与您的本地和 CI 结果保持一致。现在,使用各种 IDE 的 PyTorch 贡献者可以在编写代码时发现类型错误,并确信结果是一致的,这减少了上下文切换并更容易保持高质量的代码。VSCode 用户可以在此下载我们的 IDE 扩展。启用后,它将自动找到 PyTorch 项目中的配置文件。

先进的类型能力

Pyrefly 为 PyTorch 带来了先进的类型功能,包括对复杂类型模式的稳健支持,以及对 Python 类型规范的严格遵守。这使贡献者能够编写更安全、更具表达力的代码,同时保持高性能和流畅的开发体验。

Pyrefly 的推理能力还可以使开发者即使在缺乏显式类型注释的代码中也能检测到类型错误。这意味着遗留代码、实验性模块和快速迭代的原型可以在不需要大量预先投入注释工作的情况下,提高类型安全性。它还有助于识别可以从更明确的类型注释中受益的代码区域,从而帮助我们实现提高代码库类型覆盖率的目标。目前,PyTorch 默认未启用返回类型推断,但我们正在积极努力添加注释并修复类型问题,以便在不久的将来解锁此功能。

def foo():
    return 1
foo() + "hello" # mypy: no error, # pyrefly: error [unsupported-operation]

开始使用 Pyrefly

PyTorch 的贡献者可以通过在编辑器中安装扩展来开始使用 Pyrefly,并可以使用 lintrunner 快速轻松地进行本地类型检查:

lintrunner init
lintrunner

Helion 的贡献者也可以通过安装 IDE 扩展开始使用,并可以通过运行存储库的 lint.sh 文件进行本地类型检查。

 ./lint.sh install && ./lint.sh

Pyrefly 也已集成到我们的 CI 套件的 lint 作业中,以确保整个代码库的一致性。这确保了在本地开发中应用的相同规则在每个 PR 上都能得到执行。当你提交拉取请求时,可以通过导航到“Checks”选项卡并选择 lint 作业来找到 Pyrefly 的结果。

如果你不是 PyTorch 贡献者,但想在自己的项目中使用 Pyrefly,可以在此处获取 VSCode 扩展,或查阅 Pyrefly 文档

未来工作

切换到 Pyrefly 标志着 PyTorch 项目迈出了实用且有意义的一步。开发者们已经看到了更快速、更一致的类型检查所带来的好处,初步的推广已经帮助发现并解决了大量 Bug。这一转变精简了工作流,并为持续改进代码质量和开发者体验奠定了基础。

展望未来,我们希望随着工具的成熟,能继续看到 Pyrefly 在性能上的提升。我们也期待与 Pyrefly 团队合作,进一步改善整个代码库的类型检查。加强这个最广泛使用的 AI/ML 库之一的类型注释,将使维护者和更广泛的社区能够更自信地在生产环境中使用 PyTorch。部署一个更新、更快的类型检查器只是这一旅程的第一步。

一如既往,社区的反馈非常宝贵。我们鼓励 PyTorch 贡献者和用户在我们需要不断完善类型检查工作流时分享他们的经验、报告问题并提出改进建议。如果你有疑问或希望向 Pyrefly 团队提供反馈,可以通过 Discord 进行,或在 Pyrefly 存储库中提交 GitHub Issue 以报告错误。

最后,我们要衷心感谢 PyTorch 和 Pyrefly 团队,以及在整个过渡期间提供反馈和测试的社区成员。