Add ability to specify artifact_location when using MLFlowLogger

[rchen19]

🚀 Feature

Expose argument artifact_location of MlflowClient.create_experiment when it is called inside MLFlowLogger.experiment, this can be set with an argument when MLFlowLogger is instantiated along with tracking_uri or save_dir, to specify a custom location to save artifacts.

Motivation

When I use the save_dir argument of MLFlowLogger to specify a location, different from my current working directory, to save my mlruns/, the artifacts logged using the same MLFlowLogger still get saved under my current working directory, which is kind of messy.

[ethanwharris]

@rchen19 Thanks for the issue! I couldn't actually reproduce the behaviour you describe on master. When I change save_dir to something else with MLFlowLogger and log some artifacts (self.logger.experiment.log_artifact inside LightningModule) they get moved to the new save_dir as well. So you may find that you get the desired behaviour if you install lightning from master. I've opened #6677 to expose the artifact_location argument anyway as I think it is still useful to be able to send artifacts somewhere else. Hope that helps 😃

[rchen19]

@ethanwharris Many thanks for that info. I had been using the logger outside the lightning module, i.e. in a separate experiment script, the logger is called after the training is finished to save every file generated during training to MLFlow. Is there a way to log things inside lightning module at the very end of training?

[ethanwharris]

Hi @rchen19, sure, if you use self.logger.experiment.log_artifact inside the on_train_end hook then that will be called at the end of training, before the logger experiment is closed 😃

[rchen19]

@ethanwharris Thanks a lot, this is helpful. on_train_end is defined in ModelHooks, which I had not looked into until now. Thanks again.

[rchen19]

@ethanwharris On second look, if I would like replace the steps I do outside lightning module, after trainer.fit() is finished, I should use on_fit_end, instead of on_train_end, right? The documentation on the differences of these hooks seems to be lacking.

Another related question is that, DataHooks.teardown(stage="fit") seems can be used for the same purpose?

[ethanwharris]

@rchen19 teardown should only really be used to close open connections and things like that, rather than functions like logging. I'm not sure whether on_fit_end or on_train_end is better for your use case. Those docs could definitely use some love so please consider opening a PR once you have a better handle on the differences 😃