[MLFlowObjectStore] [2/2] Support checkpointing with MLFlow

[jerrychen109]

What does this PR do?

Follow-up PR to #2802 that integrates MLFlowObjectStore with the Trainer.

• Adds MLFlowObjectStore as a backend for RemoteUploaderDownloader
• Updates MLFlowLogger to tag MLFlow runs with the Composer run name so autoresume continues the same MLFlow run.
• Updates save folder and filenames after Event.INIT to avoid format placeholders in run logs
• Logs mlflow_experiment_id and mlflow_run_id as hyperparams

Tested manually with both unsharded and sharded checkpoints:

• Full run with save_folder set to `dbfs:/databricks/mlflow-tracking/{mlflow_experiment_id}/{mlflow_run_id}/artifacts/'
• Manually stopping run after a few checkpoints and resuming to verify that autoresume begins from latest checkpoint
• Verifying HF checkpoints uploaded at end of run
• With load_path set to the checkpoints path of a previous MLFlow run, loads the weights from that checkpoint

What issue(s) does this change relate to?

Before submitting

• Have you read the contributor guidelines?
• Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[dakinggg]

The logic looks right to me. I'm only wondering if its better to stick with not {} placeholders for the mlflow object store to avoid all of the partial format changes. I know we've gone back and forth on this, what do you think now after having implemented it this way?

[jerrychen109]

I'm personally in favor of keeping the {} placeholders now. The logic won't be that much simpler without them - we still need to replace the mlflow placeholders; each instance of partial_format would just turn into two replace() calls, which I feel is actually a little messier.

Discussed with @dakinggg over Slack, we decided that the implementation of partial_format should be safe and we're not too worried about infinite loop behavior. I restructured the code slightly to make it more clear that if the try block ever completes or throws an unexpected exception (not KeyError or IndexError), the function will return on propagate the exception.

Another benefit of keeping the curly braces + partial_format is that it should technically make format_name_with_dist and format_name_with_dist a bit more robust in the rare situation that someone actually wants to put curly braces in their path but not as one of the supported format args.