试验跟踪和可观测性

AI 运行时与 MLflow 集成，提供实验跟踪、可观察性和调试工具。 MLflow 提供了一个统一的位置来存储运行的参数、指标、GPU 系统指标、日志和产物，无论你是在笔记本中开发、使用 databricks air CLI 提交工作负载，还是调度作业。

MLflow所提供的内容

MLflow实验将相关的训练运行分组，每次运行都会提供训练配置和结果的持久记录。

• 指标和参数：记录训练损失、评估指标、学习率和超参数，然后跨次运行进行比较。
• 系统指标：在系统指标标签页查看管理运行的 GPU、CPU 和内存使用率。
• 日志视图：在运行的日志选项卡中查看由作业支持的工作负载的驱动程序输出。请参阅在 MLflow 日志选项卡中查看日志。
• 制品与模型：运行时存储模型文件、配置及其他输出。 将 Unity Catalog 卷用于大型或持久工件。
• 共享与协作：授予实验访问权限，使其他工作区用户可以查看和比较其实验运行记录。 请参阅使用 MLflow 实验来组织训练运行。
• 框架集成：使用与 Hugging Face Transformers、PyTorch Lightning 及其他支持库的 MLflow 集成。

对于由作业支持的工作负载，例如使用 databricks air run 提交的工作负载，Azure Databricks 作业运行通过记录执行状态、计算资源、尝试次数和驱动程序输出来补充 MLflow。

MLflow 运行的创建方式

你如何开始或加入MLflow运行取决于你如何开始工作负载：

你如何运行工作负载	MLflow 运行的开始方式	如何选择实验
Databricks CLI （databricks air run）	每次尝试都会自动创建一个运行。	在工作负载 YAML 中设置 experiment_name。
带有Databricks AI Runtime SDK的笔记本（.distributed()）	每次调用都会自动生成一条运行记录。	在.distributed()之前调用mlflow.set_experiment()，或设置MLFLOW_EXPERIMENT_NAME。
在附加的 Notebook 中直接运行代码	不会自动创建任何 Run。 使用 mlflow.start_run() 或自动记录。	在笔记本中调用 mlflow.set_experiment()。
已计划ai_runtime_task	该任务会自动创建一条运行记录。	在任务配置中设置 experiment。 参见 AI 运行时任务。

在你的代码启动或加入正确的运行后，使用 MLflow 跟踪 API 记录自定义训练数据。 对于支持的框架，你可以用自动日志代替单独的日志调用。 关于深度学习模式，请参见 MLflow 3深度学习工作流程。

配置MLflow跟踪时请采用以下做法：

• 将 step 设置为有意义的值，例如全局批次或纪元。 按间隔记录，而不是每批次都记录。 MLflow 限制每次运行的度量步数。 请参阅资源限制。
• 对于分布式训练，仅从排名 0 记录指标，以便每个值只被记录一次。
• 对于笔记本工作流程，可以使用绝对实验路径，比如 /Users/<username>/my-experiment 或 /Workspace/Shared/<team>/my-experiment。 把你想分享的实验放到共享工作区文件夹里。
• 将大型且持久的检查点存储在 Unity Catalog 卷中。 关于检查点模式，请参见 “提升AI运行时的训练性能和韧性”。

Note

本页示例使用MLflow 3.7及以上版本。

跟踪Databricks的CLI工作负载

当你提交一个带有 databricks air run 的工作负载时，AI Runtime 会创建一个 Azure Databricks 作业运行以提供执行详细信息，以及一个 MLflow 运行以记录实验数据。 原始提交及其重试共享一个作业运行，而每次尝试都会在同一实验中创建一个独立的MLflow运行，以便比较各次尝试。

每次托管运行时，GPU、CPU 和内存系统指标都会自动被捕获，并显示在 MLflow 系统指标 标签页上。

配置并记录到运行中

使用工作负载 YAML 选择实验，并可选配置运行名和工件位置。 以下示例展示了所需的 MLflow 字段：

experiment_name: my-training
mlflow_run_name: baseline-lr3e5
mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

关于字段需求、默认值和受支持的值，请参见 工作负载 YAML 参考。

AI 运行时通过 MLFLOW_RUN_ID 提供托管运行 ID。 要将你自己的参数、指标和伪影记录到该运行中，请用 MLflow跟踪API重新打开它：

import os

import mlflow

# Log from rank 0 only because all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
    with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
        mlflow.log_param("learning_rate", 3e-4)
        for step, loss in enumerate(training_losses):
            mlflow.log_metric("train_loss", loss, step=step)

查找运行和日志

日志可通过Databricks的CLI和MLflow运行页面获取。 使用以下命令查找工作负载、检查工作负载以及流式传输或下载日志：

databricks air list --all-status
databricks air get <job-run-id>
databricks air logs <job-run-id>
databricks air logs <job-run-id> --node 2
databricks air logs <job-run-id> --download-to ./logs/

跟踪笔记本工作负载

在笔记本中，你可以直接在附加的会话中运行代码，或通过 Databricks AI Runtime SDK 的 @distributed 修饰器运行分布式训练。

直接在笔记本里运行代码

直接在附加的笔记本中运行的代码不会自动创建 MLflow 运行记录。 选择一个实验，然后开始运行或启用自动日志：

import mlflow

mlflow.set_experiment("/Users/<username>/my-experiment")

with mlflow.start_run(run_name="baseline"):
    train()

要启用 PyTorch Lightning 自动记录，请使用以下模式：

mlflow.pytorch.autolog()

with mlflow.start_run(run_name="lightning-baseline"):
    trainer.fit(model, datamodule=datamodule)

对于 Hugging Face Transformers，配置 TrainingArguments 如下：

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="/Volumes/<catalog>/<schema>/<volume>/checkpoints",
    report_to="mlflow",
    run_name="llama7b-sft-lr3e5",
    logging_steps=50,
)

使用 Databricks AI 运行时 SDK

对于从笔记本中进行分布式训练，可以使用 Databricks AI Runtime SDK 中的 @distributed 装饰器（databricks-sdk-air）。 SDK 在 GPU 环境 6 及以上版本中在 databricks.air 命名空间下可用，或在 GPU 环境 5 及以下版本中为 serverless_gpu。

from databricks.air import distributed

@distributed()
def train():
    # Training code
    ...

train.distributed()

每次 .distributed() 调用都会自动创建一条 MLflow 运行记录。 默认实验是 /Users/{WORKSPACE_USER}/{notebook-name}。 要选择另一个实验，请在调用 .distributed()前设置：

mlflow.set_experiment("/Users/<username>/my-experiment")
train.distributed()

当你在一个活跃的运行中调用 .distributed() 时，AI 运行时会创建一个嵌套的子运行：

with mlflow.start_run() as parent_run:
    train.distributed()

要恢复之前的一次运行，请在 MLFLOW_RUN_ID 之前设置 .distributed()：

import os

os.environ["MLFLOW_RUN_ID"] = "<previous-run-id>"
train.distributed()

使用单元格输出中的运行链接或笔记本 实验 面板打开MLflow运行。

查看日志

• 笔记本输出：标准输出和训练代码产生的错误输出会显示在笔记本单元格的输出中。
• MLflow 日志：对于作业支持的笔记本工作负载（如 .distributed(remote=True) 调用），MLflow Logs 标签页显示相关作业运行的输出。

监控实时GPU资源

笔记本 GPU资源 面板显示单节点和多节点工作负载的实时GPU利用率、内存和温度。 当你需要与某次运行绑定的持久记录时，可以使用 MLflow 系统指标 标签页。

要打开该窗格，将笔记本连接到 AI 运行时，然后在右侧窗格中单击 GPU 资源。

该窗格每 10 秒轮询一次指标，最多保留 2 小时的历史记录。 点击 刷新 以获取最新值。 5分钟不操作后，窗格会暂停。 重新打开它以恢复监控。

在MLflow日志标签页查看日志

MLflow 运行页面上的 日志 标签显示了由作业支持的 AI 运行时工作负载的输出。 要打开它，请进入运行页面并点击 日志。

对于新运行，日志标签替换了之前存储在运行的工件标签中的日志文件，而该标签现已弃用。 之前创建的运行会把日志文件保存在 Artifacts 标签页上。

在 日志 标签中，你可以：

• 搜索：只显示包含搜索词的行。
• 按节点、级别和时间范围筛选：选择多节点工作负载中的单个节点，仅显示信息、警告或错误线，或缩小视图至时间范围。
• 跟踪运行工作负载：运行时，新的输出会随着到达而出现。 向上滚动以暂停跟随。
• 复制并下载：复制已加载的行，或下载所有符合当前筛选条件的行。

日志保留

日志可在 日志 标签页查看，期限为60天。

排查日志访问问题

MLflow 运行页面上的日志标签会流式传输关联的 Azure Databricks 作业运行的输出。 访问MLflow实验并不意味着能访问该作业。

如果标签显示 “您无权访问这些日志”，请请求拥有 “可管理” 权限的用户或工作区管理员，至少为您授予对该作业的 “可查看” 权限。 请参阅控制对作业的访问。

Azure Databricks 中的全局限制

请参阅资源限制。