如何在无本地文件时从MLflow加载含自定义predict的模型并解决模块缺失问题

问题核心 #

你遇到的问题本质是：code_path确实会把指定文件复制到MLflow模型的artifacts目录，但Python导入系统默认不会自动识别这个目录，导致加载时找不到自定义模块。下面是具体的解决办法：

1. 检查code_path的正确性 #

• 首先确认保存模型时code_path的路径和实际文件一致。比如你的自定义类在model.py，保存时要写：

  mlflow.pyfunc.log_model(
      artifact_path="custom_model",
      python_model=wrapper_model,
      code_path=["model.py"]  # 写实际文件路径，支持相对/绝对路径
  )
  

• 保存后可以去MLflow模型的code子目录下查看，确认model.py已被复制。如果没有，检查路径是否写错（比如误写为models.py或层级错误）。

2. 加载时手动添加代码目录到Python路径 #

MLflow不会自动把模型的code目录加入sys.path，需要手动处理：

import mlflow
import sys
import os

# 加载模型
model = mlflow.pyfunc.load_model("path/to/your/model")

# 获取模型的code目录路径
model_code_dir = os.path.join(model._model_impl._model_meta._artifact_path, "code")
# 加入sys.path
sys.path.append(model_code_dir)

# 现在可以正常使用模型
preds = model.predict(input_data, extra_param="test")

3. 更优雅的方式：在Wrapper类的load_context中处理路径 #

把路径添加逻辑封装到WrapperPythonModel的load_context方法里，加载模型时自动完成配置：

import mlflow.pyfunc
import sys
import os

class WrapperPythonModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 获取code目录路径
        code_dir = context.artifacts.get("code", os.path.dirname(context._model_path))
        sys.path.append(code_dir)
        # 正常导入自定义类
        from model import CustomForest
        # 加载预训练模型实例（根据你保存的方式调整，比如用joblib）
        import joblib
        self.model = joblib.load(os.path.join(context.artifacts["model_weights"]))

    def predict(self, context, model_input, extra_param=None):
        # 自定义predict逻辑
        base_pred = self.model.predict(model_input)
        return base_pred, extra_param

保存模型时，若有单独的权重文件，可通过artifacts参数指定：

mlflow.pyfunc.log_model(
    artifact_path="custom_model",
    python_model=wrapper_model,
    code_path=["model.py"],
    artifacts={"model_weights": "path/to/saved/model_weights.joblib"}
)

4. 检查导入语句的一致性 #

报错提示"module models.py not found"，先排查代码里的导入语句是不是写成了from models import CustomForest，但实际文件名是model.py——这种拼写错误很常见，确保导入模块名和实际文件名完全一致。

5. 长期解决方案：封装成可安装的Python包 #

如果需要频繁部署这类自定义模型，建议把model.py所在代码封装成可安装的Python包（比如my_custom_models），在训练和部署机器上都安装这个包：

• 打包成wheel或用pip install -e .安装开发版
• 这样保存模型时无需指定code_path，直接导入即可，彻底解决路径问题

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