如何在MLFlow中注册模型以外的对象（如DataFrame）

完全可行，你之前遇到的df属性丢失问题，是因为MLflow默认不会序列化PythonModel自定义的__init__参数——它只处理模型推理相关的核心逻辑，所以需要手动实现DataFrame的保存与加载逻辑，或者换更直接的方式。

方法一：重写PythonModel的序列化/反序列化逻辑 #

通过实现save_context和load_context方法，手动处理DataFrame的持久化与加载，确保属性在重新加载后保留。

import mlflow.pyfunc
import pandas as pd
import os
from mlflow.tracking import MlflowClient

class DataFrameModel(mlflow.pyfunc.PythonModel):
    def __init__(self, df=None):
        super().__init__()
        self.df = df

    def save_context(self, context, path):
        # 用parquet格式保存DataFrame（支持复杂数据类型，存储效率更高）
        df_save_path = os.path.join(path, "dataframe.parquet")
        self.df.to_parquet(df_save_path)

    def load_context(self, context):
        # 从artifacts中加载保存的DataFrame
        df_load_path = os.path.join(context.artifacts["dataframe"], "dataframe.parquet")
        self.df = pd.read_parquet(df_load_path)

# 注册DataFlow到MLflow Model Registry
if __name__ == "__main__":
    # 测试用DataFrame
    test_df = pd.DataFrame({"id": [1,2,3], "value": ["x","y","z"]})
    artifact_path = "production_df_model"

    with mlflow.start_run():
        # 日志模型并指定artifacts映射
        mlflow.pyfunc.log_model(
            artifact_path=artifact_path,
            python_model=DataFrameModel(test_df),
            artifacts={"dataframe": artifact_path}
        )
        # 注册到模型仓库
        run_id = mlflow.active_run().info.run_id
        model_uri = f"runs:/{run_id}/{artifact_path}"
        mlflow.register_model(model_uri, "ProductionDataFrame")

# 加载已注册的DataFrame
client = MlflowClient()
model_version = client.get_model_version("ProductionDataFrame", 1)
loaded_model = mlflow.pyfunc.load_model(model_version.source)
print(loaded_model.df)  # 正常访问df属性

方法二：直接将DataFrame作为Artifact保存（更轻量化） #

如果不需要封装成PythonModel，直接把DataFrame作为MLflow Artifact上传，再通过Model Registry关联管理，操作更简单。

import mlflow
import pandas as pd
from mlflow.tracking import MlflowClient

test_df = pd.DataFrame({"id": [1,2,3], "value": ["x","y","z"]})

with mlflow.start_run():
    # 保存DataFrame为parquet文件并作为artifact上传
    df_file_name = "production_data.parquet"
    test_df.to_parquet(df_file_name)
    mlflow.log_artifact(df_file_name, artifact_path="dataframes")

    # 可选：注册一个空的PythonModel来关联这个artifact，方便在Model Registry中统一管理
    mlflow.pyfunc.log_model(
        artifact_path="df_reference_model",
        python_model=mlflow.pyfunc.PythonModel(),
        artifacts={"production_df": "dataframes/production_data.parquet"}
    )
    run_id = mlflow.active_run().info.run_id
    mlflow.register_model(f"runs:/{run_id}/df_reference_model", "ProductionDataFrame")

# 加载方式
client = MlflowClient()
model_version = client.get_model_version("ProductionDataFrame", 1)
# 方式1：从模型的artifacts目录加载
loaded_df = pd.read_parquet(os.path.join(model_version.source, "artifacts/production_df"))
# 方式2：直接通过Run ID和artifact路径加载
loaded_df = pd.read_parquet(f"runs:/{run_id}/dataframes/production_data.parquet")

两种方法的适用场景 #

• 方法一适合需要将DataFrame和后续数据处理逻辑绑定的场景，加载后可以直接通过模型对象调用DataFrame和自定义方法。
• 方法二更适合单纯管理DataFrame版本的场景，操作步骤更少，轻量化。

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