Discrepancy between models loaded by xgboost and lightgbm runtimes

[ascillitoe]

Issue

The mlserver_xgboost runtime loads models as follows:

MLServer/runtimes/xgboost/mlserver_xgboost/xgboost.py

Lines 24 to 34 in 6864a2d

	def _load_sklearn_interface(model_uri: str) -> XGBModel:
	try:
	regressor = xgb.XGBRegressor()
	regressor.load_model(model_uri)
	return regressor
	except TypeError:
	# If there was an error, it's likely due to the model being a
	# classifier
	classifier = xgb.XGBClassifier()
	classifier.load_model(model_uri)
	return classifier

Whilst the mlserver_lightgbm runtime does:

MLServer/runtimes/lightgbm/mlserver_lightgbm/lightgbm.py

Line 22 in 6864a2d

self._model = lgb.Booster(model_file=model_uri)

The result is that for xgboost we end up with a sklearn API model, whereas for lightgbm we end up with the raw Booster. The latter does not have a predict_proba method (for classifiers), hence infer_output='predict_proba is not supported.

Solutions

microsoft/LightGBM#4841 suggests that lightgbm sklearn api models can (should?) be saved/loaded via joblib. We could add joblib.load support for when the model artefact's suffix is .joblib.

Alternatively, we could implement something similar to our xgboost implementation. However, converting a Booster to a sklearn api model involves accessing private attributes so might be brittle.

Related

See discussion here #1279 (comment). SKLearn api models (or models with predict_proba) are required for the white-box explainers such as TreeShap.