为何`mlflow artifacts download`比`aws s3 cp`下载速度慢很多?
MLflow制品下载慢于直接S3下载的原因及优化方案
测试场景与耗时对比
使用MLflow 1.23.1版本,将约100MB的文件分别通过MLflow和直接S3上传至S3存储桶,随后通过以下命令从命令行拉取文件:
直接S3下载命令:
time aws s3 cp s3://my-bucket/my_file.pth .
执行耗时:
real 0m2.113s
user 0m1.252s
sys 0m0.373s
MLflow制品下载命令:
time mlflow artifacts download --artifact-uri dbfs:/databricks/mlflow-tracking/.../my_file.pth
执行耗时:
real 0m4.270s
user 0m1.731s
sys 0m1.031s
耗时更长的原因
- 重复客户端身份验证:MLflow 1.23.1版本中,下载S3制品时存在多次重复的客户端验证流程,额外增加了请求往返的开销,该问题在2021年12月被修复。
- 无并行下载支持:早期MLflow未实现并行化下载逻辑,单线程处理大文件的效率远低于AWS CLI的多线程下载机制,该优化在2022年2月通过代码合并实现。
优化方法
- 升级MLflow版本:将MLflow升级到包含上述两个修复的新版本(建议升级至1.26.0及以上版本,其中身份验证修复在1.26.0中合并,并行下载优化在1.28.0中合并),直接获得官方修复的性能提升。
- 直接使用AWS CLI下载:若已知制品对应的S3路径,直接使用
aws s3 cp命令下载,绕开MLflow的封装逻辑,减少额外开销。 - 使用S3直接URI替代DBFS路径:将
artifact-uri替换为直接的S3路径(s3://...),避免MLflow解析DBFS到S3路径的中间环节,减少网络请求和处理时间。 - 配置MLflow S3客户端参数:升级后,可通过配置
MLFLOW_S3_DOWNLOAD_PARALLELISM环境变量开启并行下载,或调整MLFLOW_S3_RETRY_COUNT等参数优化请求重试逻辑,提升下载效率。
内容的提问来源于stack exchange,提问作者tiberius
相关产品推荐
相关产品推荐

