Now you can log the model using this parameter:

mlflow.sklearn.log_model(
    ...,  # the usual params
    pyfunc_predict_fn="predict_proba"
)

 which will return probabilities for the first class apparently when using the model for inference (e.g. when loading it using mlflow.pyfunc.spark_udf() ).