Update to the 2.x API

Since MLFlow 2.x has been out for a while now and is pretty stable an upgrade to the 2.x API should be on the horizon.

Note that currently it mostly works fine as is, with a few rough edges that I'm investigating.