Using MLFlow with CrewAI

Anyone knows how to use MLFlow to trace inputs / outputs, log time taken, etc?
Documentation in CrewAI Docs has been quite sparse.

I have put in the following as per the logs, but am not sure how to log individual agents.

import mlflow

mlflow.crewai.autolog()
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("CrewAI")

Noticed the above 3 lines of code works well when using a single python script with agents, tasks, crew, and the kickoff altogether in a file.

However, I cant seem to get ML Flow to work when using the CrewAI YAML folder structure, or for more complex codes where the various crewAI components are split into several files.