Conversation
There was a problem hiding this comment.
Code Review
This pull request correctly adds the logic to finish MLflow runs by calling mlflow.end_run(). I've identified a couple of areas for improvement. First, the cleanup logic is placed in the __del__ method, which can be unreliable in Python; I recommend refactoring the Tracking class to use a context manager for deterministic cleanup. Second, the _MlflowLoggingAdapter class now has a duplicated import mlflow statement, which could be avoided by a small refactoring of the class. My detailed comments provide more information and suggestions.
| def finish(self): | ||
| import mlflow | ||
| mlflow.end_run() |
There was a problem hiding this comment.
The import mlflow statement is now present in both the log method and this new finish method. To avoid this duplication and improve maintainability, it would be better to refactor the _MlflowLoggingAdapter class. You could add an __init__ method that imports mlflow and stores it as an instance attribute (e.g., self.mlflow). Then, both log and finish methods can use self.mlflow without needing to import it themselves. This pattern is already used in _TensorboardAdapter in this same file.
There was a problem hiding this comment.
In this PR, I'm following the existing pattern to avoid introducing too many different patterns. I'd be happy to refactor both _MlflowLoggingAdapter and _TensorboardAdapter if we think it's the way to go.
Edit: oh, my bad, I read too quickly, _TensorboardAdapter is already importing tensorboard in the _init_.
I add import mlflow in finish to reduce the diff to a minimum, but happy to move it to an _init_ function if we think it's the way to go.
This PR calls the `end_run` MLflow method to finish the run at the end.
This PR calls the
end_runMLflow method to finish the run at the end.