Skip to content

[fix] finish mlflow run#243

Merged
SumanthRH merged 2 commits intoNovaSky-AI:mainfrom
DataDog:etnbrd/end-mlflow-run
Sep 5, 2025
Merged

[fix] finish mlflow run#243
SumanthRH merged 2 commits intoNovaSky-AI:mainfrom
DataDog:etnbrd/end-mlflow-run

Conversation

@etnbrd
Copy link
Contributor

@etnbrd etnbrd commented Sep 5, 2025

This PR calls the end_run MLflow method to finish the run at the end.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +185 to +187
def finish(self):
import mlflow
mlflow.end_run()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Copy link
Contributor Author

@etnbrd etnbrd Sep 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@SumanthRH SumanthRH merged commit e74c484 into NovaSky-AI:main Sep 5, 2025
3 checks passed
dzorlu referenced this pull request in fleet-ai/SkyRL Feb 4, 2026
This PR calls the `end_run` MLflow method to finish the run at the end.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants