如何在Databricks Notebook中结合Metaflow与MLflow运行工作流？

问题根源 #

Databricks环境中，MLflow默认通过内部API完成认证，但Metaflow调用MLflow时会尝试读取本地CLI配置文件（~/.databrickscfg），而Notebook环境默认没有这个配置，因此触发认证错误。

解决步骤 #

1. 直接在代码中配置MLflow追踪环境 #

不需要依赖CLI配置，在Metaflow脚本开头添加MLflow的Databricks追踪配置：

import mlflow

# 设定MLflow追踪目标为当前Databricks工作区
mlflow.set_tracking_uri("databricks")
# 可选：指定具体实验路径（不指定则使用默认实验）
mlflow.set_experiment("/Users/你的邮箱@xxx.com/目标实验名")

2. 在Metaflow步骤内正确调用MLflow #

将MLflow的日志逻辑放在Metaflow的@step装饰函数内部，避免初始化阶段触发认证问题：

from metaflow import FlowSpec, step

class MLflowMetaflowFlow(FlowSpec):
    @step
    def start(self):
        import mlflow
        # 启动MLflow运行上下文
        with mlflow.start_run():
            # 日志参数
            mlflow.log_param("learning_rate", 0.01)
            # 日志指标
            mlflow.log_metric("accuracy", 0.95)
        self.next(self.end)

    @step
    def end(self):
        print("工作流执行完成")

if __name__ == "__main__":
    MLflowMetaflowFlow()

3. Notebook中运行脚本的可靠方式 #

放弃!python命令，改用以下两种方式之一：

• 方式一：将脚本代码直接放在Notebook单元格中执行，共享Notebook的环境认证信息
• 方式二：若脚本存储在DBFS，使用Databricks魔法命令运行：

  %run /dbfs/path/to/your/script.py
  

如果必须用!python，需提前注入环境变量传递认证信息：

!export DATABRICKS_HOST=$DATABRICKS_HOST && export DATABRICKS_TOKEN=$DATABRICKS_TOKEN && python your_script.py

4. 验证权限 #

确保运行Notebook的用户具备目标MLflow实验的读写权限，以及Metaflow所需的存储资源（如DBFS、S3）权限。

关键注意事项 #

• 不要依赖本地~/.databrickscfg文件，Databricks Notebook已通过环境变量注入认证信息，mlflow.set_tracking_uri("databricks")会自动读取这些信息完成认证
• 务必在@step函数内部导入MLflow，避免Metaflow初始化阶段提前触发MLflow的认证逻辑

内容的提问来源于stack exchange，提问作者Daniel