That's what I ended up doing. In order to do that however, I need to create clusters with different Databricks runtimes and choose the latest runtime that had Python 3.8.*. It was a lot more work than I anticipated.
I think the Databricks docs should be a lot clearer.