Databricks中MLflow嵌套运行在自定义实验位置无法分组问题

在Databricks中使用MLflow嵌套运行实现XGBoost超参数优化时，遇到以下异常：

• 默认实验位置下，父运行xgboost_models可正常展开显示子运行，分组逻辑正常。
• 切换到自定义实验位置后，父运行xgboost_models_2无任何子运行，所有子运行均出现在默认实验位置，且GUI无法识别父子关系（尽管子运行的Parent ID已正确指向父运行ID）。

自定义实验创建代码：

EXPERIMENT_NAME = '/Users/dxxxx@redacted.com/MLflow_experiments/dxxxx_minimal_MLflow'

# Get the experiment ID if it exists, or create a new one
experiment_id = mlflow.get_experiment_by_name(EXPERIMENT_NAME)

if experiment_id is None:
    # If the experiment does not exist, create it
    experiment_id = mlflow.create_experiment(EXPERIMENT_NAME)
else:
    # If the experiment exists, get its ID
    experiment_id = experiment_id.experiment_id

修改后的父运行代码：

with mlflow.start_run(experiment_id=experiment_id, run_name='xgboost_models_2'):
    best_params = fmin(
        fn=train_model, 
        space=search_space, 
        algo=tpe.suggest, 
        max_evals=96,
        trials=spark_trials,
    )

尝试在子运行的mlflow.start_run()中添加experiment_id和parent_run_id参数，问题仍未解决。

1. 问题解决方法 #

核心原因是SparkTrials在分布式环境下无法自动继承Driver端的MLflow实验上下文，导致子运行默认使用默认实验。需通过以下步骤修复：

步骤1：全局设置实验上下文 #

在启动父运行前，先通过mlflow.set_experiment全局指定实验ID，确保所有后续运行（包括分布式子任务）默认使用该实验：

# 全局设置实验，替代父运行中的experiment_id参数
mlflow.set_experiment(experiment_id=experiment_id)

with mlflow.start_run(run_name='xgboost_models_2'):
    best_params = fmin(
        fn=train_model, 
        space=search_space, 
        algo=tpe.suggest, 
        max_evals=96,
        trials=spark_trials,
    )

步骤2：确保子运行启用嵌套模式 #

在train_model函数中，启动子运行时必须添加nested=True参数，同时无需手动指定experiment_id（会自动继承全局上下文）：

def train_model(params):
    # 启用嵌套运行，自动继承父运行的实验ID和父ID
    with mlflow.start_run(nested=True):
        mlflow.log_params(params)
        # 后续训练、评估、日志代码
        # ...

步骤3：验证SparkTrials上下文传递 #

如果仍有问题，可在train_model函数中显式打印实验ID和父运行ID，确认上下文是否正确传递：

def train_model(params):
    with mlflow.start_run(nested=True):
        print(f"Current experiment ID: {mlflow.active_run().info.experiment_id}")
        print(f"Parent run ID: {mlflow.active_run().info.parent_run_id}")
        # ...

2. 问题范围说明 #

这不是MLflow的通用问题，而是Databricks平台结合SparkTrials分布式运行时的特定问题：

• 本地单机运行MLflow嵌套时，子运行会自动继承父运行的实验上下文，父子关系可正常显示。
• 在Databricks的Spark分布式环境中，SparkTrials的每个子任务运行在独立的Executor节点上，无法自动继承Driver端的MLflow实验配置，导致子运行默认落入默认实验。即使子运行的Parent ID正确，由于父子运行分属不同实验，GUI无法识别并展示分组关系。

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