如何在mlflow.pyfunc.log_model中指定必要依赖或排除冗余Python包？

问题原因 #

你的SetfitModelWrapper类定义在my_app模块下，MLflow用cloudpickle序列化模型时，会保留这个类的模块路径信息。生产环境加载模型时，会尝试导入my_app模块，但环境里没装就会触发ModuleNotFoundError。

解决方案 #

1. 将模型Wrapper移到独立的无依赖模块

把SetfitModelWrapper从my_app里抽出来，放到一个只依赖MLflow和SetFit的独立文件中，避免序列化时带上my_app的依赖。

步骤：

• 新建文件model_wrappers.py，内容如下：

import mlflow.pyfunc
from setfit import SetFitModel

class SetfitModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model: SetFitModel):
        self.model = model

    def predict(self, context, model_input):
        # 补全你的预测逻辑
        return self.model.predict(model_input)

• 在my_app中导入这个Wrapper来包装模型：

from model_wrappers import SetfitModelWrapper
from setfit import SetFitModel

# 加载你的SetFitModel
model = SetFitModel.from_pretrained("your-model-path")
wrapped_model = SetfitModelWrapper(model=model)

# 日志模型
mlflow.pyfunc.log_model(
    artifact_path='./artifact',
    python_model=wrapped_model
)

2. 手动指定打包的代码和依赖环境

用mlflow.pyfunc.log_model的code_paths参数指定需要打包的代码文件，同时用conda_env参数定义模型运行所需的最小依赖，避免包含整个my_app。

示例代码：

import mlflow.pyfunc
from setfit import SetFitModel

class SetfitModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model: SetFitModel):
        self.model = model

    def predict(self, context, model_input):
        return self.model.predict(model_input)

model = SetFitModel.from_pretrained("your-model-path")
wrapped_model = SetfitModelWrapper(model=model)

# 定义最小conda环境依赖
conda_env = {
    "name": "setfit-model-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.11",
        "pip",
        {
            "pip": [
                "mlflow>=2.0",
                "setfit>=1.0",
                "transformers>=4.0"
            ]
        }
    ]
}

# 日志模型时指定code_paths和conda_env
mlflow.pyfunc.log_model(
    artifact_path='./artifact',
    python_model=wrapped_model,
    code_paths=["./path/to/your/wrapper/file.py"],  # 只打包Wrapper所在文件
    conda_env=conda_env
)

3. 自定义序列化逻辑（不推荐）

通过修改Wrapper类的__reduce__方法，让cloudpickle序列化时不依赖my_app模块，但这种方式容易出错，仅作为备选：

class SetfitModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model: SetFitModel):
        self.model = model

    def predict(self, context, model_input):
        return self.model.predict(model_input)

    def __reduce__(self):
        # 返回类本身和初始化参数，避免依赖原模块路径
        return (SetfitModelWrapper, (self.model,))

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