跳到主要内容

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 时将其替换为 mlflow-go。这将启动一个 Go Web 服务器,其 API 与 MLflow REST API 完全一致。

在撰写本文时,用户需要将数据库连接字符串作为 '--backend-store-uri' 参数传递,因为 Go 实现目前仅支持数据库。由于使用 mlflow-go 的目的是为了其性能,因此使用一个实际的生产级数据库而非文件系统是合理的。

总结思考

尽管我们在 MLflow-go 中构建的后端对于某些大数据应用来说性能更优,但这并不意味着它会取代现有的 Python 实现。由于 Databricks 内部的 MLflow 维护者并非专业的 Go 语言开发者,mlflow-go 在被广泛采用之前,必须发展出自己的用户群和活跃贡献者社区。

我们发现 G-Research 非常喜欢 MLflow,但需要其性能更好,特别是对于大型数据集。我们预计其他组织,尤其是那些拥有大量数据处理需求的组织,也遇到了类似的问题。我们非常希望了解我们的努力是否让其他人受益,因此请在 GitHubSlack 上告知我们您如何使用 MLflow,以及 FastTrackML 对您来说是否更快。