如何在Databricks中程序化获取MLflow已记录工件的URL？

问题背景 #

我正在运行机器学习流水线，使用MLflow记录实验信息，参考Databricks官方MLflow追踪教程，执行代码如下：

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

with mlflow.start_run():
  n_estimators = 100
  max_depth = 6
  max_features = 3
  # 创建并训练模型
  rf = RandomForestRegressor(n_estimators = n_estimators, max_depth = max_depth, max_features = max_features)
  rf.fit(X_train, y_train)
  # 生成预测结果
  predictions = rf.predict(X_test)
  
  # 记录参数
  mlflow.log_param("num_trees", n_estimators)
  mlflow.log_param("maxdepth", max_depth)
  mlflow.log_param("max_feat", max_features)
  
  # 记录模型
  mlflow.sklearn.log_model(rf, "random-forest-model")
  
  # 计算评估指标
  mse = mean_squared_error(y_test, predictions)
    
  # 记录指标
  mlflow.log_metric("mse", mse)

在Databricks笔记本中运行后，会显示状态提示：

(1) MLflow run
Logged 1 run to an experiment in MLflow. Learn more

我可以点击"1 run"查看记录的信息，但希望程序化获取该Run的完整UI链接，链接格式为：

https://mycompany-dev.cloud.databricks.com/?o=<ID_1>#mlflow/experiments/<ID_2>/runs/<ID_3>

其中<ID_1>是Databricks工作区ID，<ID_2>是experiment_id，<ID_3>是run_id。我不想通过硬编码拼接的方式构建URL，需要程序化生成的方案。

解决方案 #

1. 在Run上下文内直接生成链接 #

利用Databricks内置工具获取工作区ID，结合MLflow Run对象的信息动态构建链接：

import mlflow
import os

with mlflow.start_run() as run:
    # 你的模型训练、日志代码...
    
    # 从Databricks环境获取工作区ID
    workspace_id = dbutils.notebook.entry_point.getDbutils().notebook().getContext().tags().get("orgId").get()
    
    # 动态获取Databricks基础域名（避免硬编码）
    databricks_base_url = os.getenv("DATABRICKS_HOST")
    
    # 构建完整Run UI链接
    run_ui_url = f"{databricks_base_url}/?o={workspace_id}#mlflow/experiments/{run.info.experiment_id}/runs/{run.info.run_id}"
    
    # 输出或记录链接
    print("MLflow Run UI链接:", run_ui_url)

2. 事后通过Run ID查询生成链接 #

如果需要在Run结束后获取链接，可以使用MLflow Client查询Run信息：

from mlflow.tracking import MlflowClient
import os

client = MlflowClient()
# 替换为目标Run ID
target_run_id = "<your-run-id>"
run = client.get_run(target_run_id)

# 获取工作区ID和基础域名
workspace_id = dbutils.notebook.entry_point.getDbutils().notebook().getContext().tags().get("orgId").get()
databricks_base_url = os.getenv("DATABRICKS_HOST")

# 构建链接
run_ui_url = f"{databricks_base_url}/?o={workspace_id}#mlflow/experiments/{run.info.experiment_id}/runs/{run.info.run_id}"
print(run_ui_url)

关键说明 #

• DATABRICKS_HOST是Databricks环境内置的环境变量，自动指向当前工作区的域名，无需手动配置
• orgId标签对应Databricks工作区ID，通过dbutils工具可以安全获取，无需硬编码

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