在Databricks Model Serving中使用私有Python包的问题排查

一、DBFS中whl包无法找到的解决方案 #

Model Serving的构建容器不会自动挂载DBFS到本地文件系统，因此直接使用/dbfs/...或dbfs:/...路径无法被pip识别。正确做法是将私有包whl文件作为模型artifact一起上传：

步骤1：将whl文件复制到模型输出目录 #

在记录模型前，把DBFS上的whl文件复制到MLflow模型的输出子目录中：

import os
import shutil

# 源whl文件的DBFS路径
whl_dbfs_path = "/dbfs/FileStore/tables/private_package-0.1.10-py3-none-any.whl"
# 模型输出目录下的whl存放子目录
model_whl_dir = os.path.join("hello-world", "whls")
os.makedirs(model_whl_dir, exist_ok=True)

# 复制文件到模型目录
shutil.copy(whl_dbfs_path, model_whl_dir)

步骤2：用相对路径记录模型 #

在mlflow.pyfunc.log_model的pip_requirements中使用相对于模型根目录的路径：

mlflow.pyfunc.log_model(
    "hello-world",
    python_model=model,
    registered_model_name="hello-world",
    pip_requirements=[
        "./whls/private_package-0.1.10-py3-none-any.whl"
    ],
)

验证方法 #

在MLflow UI中打开已记录的模型，查看hello-world/whls目录下是否存在目标whl文件，确认artifact上传成功。

二、私有PyPI Token安全注入方案 #

无需明文存储token，可通过Databricks Secrets结合环境变量注入：

步骤1：创建Secret与Token #

1. 在Databricks Workspace中创建一个secret scope（如mlflow-private-packages）。
2. 在该scope中添加secret（如private_pypi_token），值为你的私有PyPI访问token。

步骤2：记录模型时注入环境变量 #

使用environment_variables参数引用secret，在pip_requirements中用环境变量替换token：

mlflow.pyfunc.log_model(
    "hello-world",
    python_model=model,
    registered_model_name="hello-world",
    pip_requirements=[
        "--index-url https://__token__:${PRIVATE_PYPI_TOKEN}@your-private-pypi-domain/simple/",
        "private-package==0.1.10"
    ],
    env_manager="virtualenv",  # 推荐使用virtualenv以更好支持环境变量
    environment_variables={
        "PRIVATE_PYPI_TOKEN": "{{secrets/mlflow-private-packages/private_pypi_token}}"
    }
)

三、调试建议 #

1. 权限检查：确认部署Model Serving的身份（用户或服务主体）拥有模型artifact的读取权限，以及对应secret scope的读取权限。
2. 路径验证：在notebook中打印模型输出目录的结构，确认whl文件复制路径正确。
3. 构建日志分析：查看Model Serving的构建日志，若仍有路径错误，检查artifact是否完整上传，相对路径是否匹配模型目录结构。

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