跳到主要内容

MLflow Go

·3 分钟阅读
Florian Verdonck
G-Research 开源开发者

MLflow 项目是机器学习社区的基石。它高度灵活,让机器学习工作流更具可复现性、可管理性和协作性。

MLflow 的一个关键部分是实验跟踪,即用于记录和查询实验的系统。开发者可以记录参数、指标、输入和工件,以跟踪不同的实验运行,将它们存储到数据库中,然后比较结果。

MLflow 中的这个实验跟踪功能是用 Python 编写的,在处理海量数据时存在性能限制。在 G-Research,两个月内产生 10TB 的数据量并不少见。在这种情况下,我们的研究人员可能会因为系统执行操作所需的时间过长而感到沮丧。

为了解决这个问题,我们拆解了 MLflow 的各个组件,并很快看到了系统可以改进的地方——它的后端。对于我们必须执行的海量操作而言,写入用 Python 编写的数据库是行不通的。我们开始设想用另一个后端替换实验跟踪系统。这一思路演变成了名为 FastTrackML 的项目。

Fasttrack ML

Fasttrack ML 对于 G-Research 来说是一次巨大的成功,因为它用 Go 语言实现后速度显著提升,并且能够处理海量数据。选择 Go 是因为它的并发模型基于轻量级的 goroutine 和 channel,这使得它在处理多个并发任务(例如,记录实验数据、处理用户请求)时效率极高。Go 是一种编译型语言,能生成本地机器码,因此与解释型的 Python 相比,执行速度更快。Go 是静态类型的,与 Python 的动态类型相比,内存使用效率更高。

在某种程度上,FastTrackML 的成功释放了 MLflow 在 G-Research 内部的使用潜力。它让量化研究员能够在合理的时间内处理他们的数据集和实验跟踪。

回归 MLflow

然后,G-Research 的开源团队研究了如何将从 FastTrackML 中学到的经验融入 MLflow 本身——mlflow-go 项目就是朝这个方向迈出的第一步。它是一个 Python 包,旨在成为 MLflow 实验跟踪服务器的直接替代品。一旦 mlflow-go 包被证明是成熟的技术,我们希望它能被 MLflow 吸收。

开始使用 mlflow-go 非常简单,只需安装 MLflow,运行 pip install mlflow-go-backend,然后在命令行运行时用 mlflow-go 替换 mlflow 即可。这将启动一个 Go Web 服务器,其 API 与 mlflow REST API 完全相同。

在撰写本文时,用户需要通过“--backend-store-uri”参数传递一个数据库连接字符串,因为 Go 的实现目前仅支持数据库。由于使用 mlflow-go 的初衷是为了其性能,因此使用一个真正适合生产环境的数据库而不是文件系统是合情合理的。

结束语

虽然我们在 MLflow-go 中构建的后端对于某些大数据应用性能更佳,但这并不一定意味着它会取代现有的 Python 实现。由于 Databricks 内部的 MLflow 维护者并非专门的 Golang 开发者,MLflow-go 必须发展出自己的用户群和活跃的贡献者社区,然后才可能被广泛采用。

我们发现 G-Research 非常喜欢 MLflow,但需要它性能更强,特别是在处理大型数据集时。我们预计其他组织,尤其是那些有海量数据处理需求的组织,也遇到了类似的问题。我们很想知道我们的努力是否对他人有所帮助,所以请在 GitHubSlack 上告诉我们您是如何使用 MLflow 的,以及 FastTrackML 对您来说是否更快。