Databricks中使用MLflow加载RoBERTa模型时出现OSError求助

问题背景 #

在Databricks环境中使用MLflow与RoBERTa Transformers训练模型，模型注册成功，但调用测试时触发如下错误：

OSError: We couldn't connect to 'https://huggingface.co/' to load this model and it looks like dbfs:/databricks/mlflow-tracking/3638851642935524/cd4eae6034684211933b97b178e5f062/artifacts/checkpoint-36132/artifacts/checkpoint-36132 is not the path to a directory containing a config.json file.

手动检查指定DBFS路径，确认config.json及其他模型文件存在，但仍提示：

Couldn't load model information due to an error.

可能原因 #

1. 模型保存时出现路径嵌套：错误路径显示checkpoint-36132/artifacts/checkpoint-36132，说明保存逻辑可能重复嵌套了模型目录，导致MLflow加载时找不到正确的文件位置。
2. MLflow与Transformers路径解析不兼容：不同版本的MLflow和Transformers对DBFS路径的处理逻辑存在差异，导致路径解析错误。
3. DBFS文件权限问题：集群运行账号无目标路径的读取权限，无法访问模型文件。

解决步骤 #

• 修正模型保存逻辑：检查训练代码中MLflow保存模型的部分，避免嵌套保存。确保模型的config.json、pytorch_model.bin等核心文件直接存放在MLflow artifacts的根目录下，而非多层嵌套的子目录中。例如，若训练时用了Hugging Face的Trainer，需确认save_model的路径设置正确，不要在checkpoint目录下再次创建artifacts子目录。
• 手动指定加载路径：调用模型时，先将模型下载到本地临时目录，再从该目录加载，绕开MLflow自动生成的嵌套路径问题：

  import mlflow.pyfunc
  # 下载模型到DBFS临时目录
  model_path = mlflow.pyfunc.load_model("models:/your_model_name/your_version", dst_path="/dbfs/tmp/roberta_model")
  # 从临时目录加载模型
  from transformers import AutoModelForSequenceClassification, AutoTokenizer
  model = AutoModelForSequenceClassification.from_pretrained("/dbfs/tmp/roberta_model")
  tokenizer = AutoTokenizer.from_pretrained("/dbfs/tmp/roberta_model")
  

• 验证DBFS权限与文件存在性：在Databricks notebook中执行以下命令，确认目标路径下的文件可被读取：

  dbutils.fs.ls("dbfs:/databricks/mlflow-tracking/3638851642935524/cd4eae6034684211933b97b178e5f062/artifacts/checkpoint-36132")
  

  若返回文件列表包含config.json，则权限正常；若报错，需调整路径的ACL权限。
• 对齐依赖版本：确保训练和推理环境的Transformers、MLflow版本完全一致。例如，训练时用transformers==4.30.2和mlflow==2.5.0，推理环境也需安装相同版本。
• 重新注册模型：若上述步骤无效，重新训练模型，在保存时直接通过MLflow的transformers.log_model方法保存，避免自定义保存逻辑导致的路径问题：

  import mlflow.transformers
  mlflow.transformers.log_model(
    transformers_model=model,
    tokenizer=tokenizer,
    artifact_path="roberta_model",
    registered_model_name="your_model_name"
  )
  

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