Databricks加载Unity Catalog的LightGBM模型丢失内置属性求助

方案1：自定义PyFunc包装类保留原模型所有属性 #

直接用mlflow.spark.log_model会将模型封装为Spark MLlib PipelineModel，丢失SynapseML LightGBM模型的原生方法。通过自定义PyFunc包装类，可以完整保留原模型的属性与方法，同时兼容MLflow的模型注册逻辑。

import mlflow
import mlflow.pyfunc
from synapse.ml.lightgbm import LightGBMClassificationModel

class LightGBMPyFuncWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 加载原始SynapseML LightGBM模型
        self.model = LightGBMClassificationModel.load(context.artifacts["model_path"])
    
    def predict(self, context, model_input):
        # 实现PyFunc标准predict接口，复用原模型的transform方法
        return self.model.transform(model_input)
    
    # 自动转发所有未定义的属性/方法到原模型
    def __getattr__(self, name):
        return getattr(self.model, name)

# 假设已训练好的原模型
original_model = LightGBMClassificationModel.load("/path/to/your/trained/model")

# 日志模型到Unity Catalog
with mlflow.start_run():
    # 临时保存原模型到DBFS
    temp_model_path = "/dbfs/temp/lgbm_original_model"
    original_model.save(temp_model_path)
    
    mlflow.pyfunc.log_model(
        artifact_path="lgbm_model",
        python_model=LightGBMPyFuncWrapper(),
        artifacts={"model_path": temp_model_path},
        registered_model_name="your_catalog.your_schema.your_lgbm_model"
    )

# 加载模型并验证属性
loaded_model = mlflow.pyfunc.load_model("models:/your_catalog.your_schema.your_lgbm_model/latest")
print(loaded_model.getFeatureImportances())  # 可正常调用原模型方法

方案2：提取原生LightGBM Booster单独保存 #

SynapseML的LightGBM模型内置原生LightGBM Booster对象，提取后用mlflow.lightgbm模块日志，能完整保留原生模型的所有方法，同时可按需搭配SynapseML模型使用。

import mlflow
import mlflow.lightgbm
from synapse.ml.lightgbm import LightGBMClassificationModel

# 训练好的SynapseML模型
synapse_lgbm_model = LightGBMClassificationModel.load("/path/to/your/trained/model")
# 提取原生Booster
native_booster = synapse_lgbm_model.getLightGBMBooster()

with mlflow.start_run():
    # 日志原生LightGBM模型到Unity Catalog
    mlflow.lightgbm.log_model(
        native_booster,
        artifact_path="native_lgbm",
        registered_model_name="your_catalog.your_schema.native_lgbm_model"
    )
    # 可选：同时日志SynapseML模型用于Spark批处理场景
    mlflow.spark.log_model(
        synapse_lgbm_model,
        artifact_path="synapse_lgbm",
        registered_model_name="your_catalog.your_schema.synapse_lgbm_model"
    )

# 加载原生模型调用原生方法
loaded_native_model = mlflow.lightgbm.load_model("models:/your_catalog.your_schema.native_lgbm_model/latest")
print(loaded_native_model.feature_importance())

# 如需Spark场景，可将原生Booster关联回SynapseML模型
loaded_synapse_model = mlflow.spark.load_model("models:/your_catalog.your_schema.synapse_lgbm_model/latest")
loaded_synapse_model.setLightGBMBooster(loaded_native_model)
print(loaded_synapse_model.getFeatureImportances())

方案3：额外保存完整模型文件用于加载 #

如果需要保持Spark模型的兼容性，同时要获取原模型的完整属性，可以在日志Spark模型时，额外保存一份完整的原模型文件到MLflow artifacts，加载时直接从该文件恢复原模型。

import mlflow
from synapse.ml.lightgbm import LightGBMClassificationModel

# 训练好的原模型
original_model = LightGBMClassificationModel.load("/path/to/your/trained/model")

with mlflow.start_run():
    # 日志Spark模型到Unity Catalog
    mlflow.spark.log_model(
        original_model,
        artifact_path="spark_lgbm",
        registered_model_name="your_catalog.your_schema.spark_lgbm_model"
    )
    # 额外保存完整原模型文件到artifact
    full_model_path = "/dbfs/temp/full_lgbm_model"
    original_model.save(full_model_path)
    mlflow.log_artifact(full_model_path, artifact_path="full_model")

# 从artifact加载完整原模型
run_id = mlflow.active_run().info.run_id
loaded_full_model = LightGBMClassificationModel.load(f"{mlflow.get_run(run_id).info.artifact_uri}/full_model")
print(loaded_full_model.getFeatureImportances())

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