如何在Azure Databricks机器学习Notebook中使用自定义Docker镜像

嘿，我来给你拆解下完整流程，自己在项目里实操过好几次，踩过的坑也一并给你提个醒~

第一步：准备符合Databricks要求的自定义镜像 #

Databricks的集群依赖官方Runtime里的核心组件（比如Spark、Databricks服务端通信工具），所以必须基于官方Databricks Runtime镜像来构建自定义镜像，不能从零开始做。

1. 编写Dockerfile示例 #

比如我要做一个带TensorFlow和PyTorch Lightning的ML镜像，Dockerfile可以这么写：

# 选一个稳定的Databricks Runtime基础镜像，比如13.3-LTS版本
FROM databricksruntime/standard:13.3-LTS

# 安装你的机器学习依赖，按需调整版本
RUN pip install tensorflow==2.15.0 pytorch-lightning==2.1.3 --no-cache-dir

# 如果有自定义工具脚本，复制到镜像里（比如自己封装的ML工具）
COPY ./my_ml_utils /databricks/my_ml_utils

小贴士：用--no-cache-dir可以减小镜像体积，避免缓存冗余。

2. 构建并推送镜像到Azure容器注册表（ACR） #

Databricks集群需要能拉取到你的镜像，所以得把镜像推送到ACR（或者其他可公网访问的私有容器仓库）：

# 先登录你的ACR
az acr login --name your-acr-registry-name

# 构建镜像，标签格式要符合ACR要求
docker build -t your-acr-registry-name.azurecr.io/databricks-custom-ml:v1 .

# 推送到ACR
docker push your-acr-registry-name.azurecr.io/databricks-custom-ml:v1

3. 给ACR设置访问权限 #

确保Databricks集群能拉取镜像：

• 如果你用的是Azure AD集成，在ACR的「访问控制（IAM）」里，给Databricks的服务主体添加「AcrPull」角色
• 或者临时用用户名密码（不推荐生产环境），在ACR的「设置→访问密钥」里开启管理员用户，记录用户名和密码

第二步：配置Databricks集群使用自定义镜像 #

1. 登录你的Azure Databricks workspace，进入「Compute」页面，点击「Create Cluster」
2. 找到「Container」选项卡（如果看不到，需要Workspace管理员在「Admin Settings→Workspace Settings」里开启「Custom containers」功能）
3. 在「Container image」输入框里填你刚才推送到ACR的镜像地址，比如your-acr-registry-name.azurecr.io/databricks-custom-ml:v1
4. 配置镜像拉取凭据：
   • 选「Azure Container Registry」，然后选择你的ACR实例（自动集成权限）
   • 或者手动输入ACR的用户名和密码（适合测试场景）
5. 其他集群配置（节点类型、数量、Runtime版本）：注意Runtime版本要和你基础镜像的版本一致（比如基础镜像是13.3-LTS，这里也选13.3-LTS）
6. 点击「Create Cluster」，等待集群初始化完成（如果镜像大，初始化时间会久一点，耐心等）

第三步：验证自定义镜像是否生效 #

集群启动后，打开一个Notebook，选择这个集群作为运行环境，跑几行代码验证：

import tensorflow as tf
import pytorch_lightning as pl

# 打印依赖版本，确认是否安装成功
print(f"TensorFlow版本: {tf.__version__}")
print(f"PyTorch Lightning版本: {pl.__version__}")

# 如果有自定义脚本，也可以导入测试
from my_ml_utils import custom_train_func
custom_train_func()

如果能正常输出版本号和执行自定义函数，说明镜像配置成功了！

一些避坑小贴士 #

• 不要修改官方镜像里的Databricks核心组件（比如/databricks下的系统文件），不然集群可能启动失败
• 镜像版本要和Runtime严格对应，比如基础镜像是13.3-LTS，集群Runtime就不能选14.0
• 生产环境要给镜像打版本标签（比如v1、v2），方便回滚和管理
• 优化镜像大小：用多阶段构建、清理安装缓存、只安装必要依赖，能缩短集群启动时间

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