Hey, at my team we are faced with the same problem.

Given that MLflow usage is part of the "Scalable Machine Learning with Apache Spark" course I got the impression that Spark ML models would work well with MLflow and Databricks model logging, but it's clearly not the case. I went back to the course material and, sure enough, when explaining the model registry it switches to using sklearn models. I gotta say, I feel a bit betrayed...

The solution given by @Retired_mod partially works but it is a hack and, as @MohsenJ said, it doesn't really cover the signature of the output. Is there a proper/more complete solution?