feat: MLFlow Integration for experiment tracking #534

therealnaveenkamal · 2025-06-21T17:42:25Z

What does this PR do ?

Add MLflow integration for experiment tracking and model management in NeMo RL

Issues

Closes #514 - Support for MLflow
This PR addresses the community request for MLflow support alongside existing TensorBoard and Weights & Biases logging options.

Usage

# Install MLflow
uv pip install mlflow

# Run SFT training with MLflow logging
uv run python examples/run_sft.py --config examples/configs/sft_mlflow.yaml

# Start MLflow UI to view results
mlflow ui --host 0.0.0.0 --port 5000

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

This PR adds comprehensive MLflow integration including:

MLflowLogger class with support for metrics, hyperparameters, and plot logging
Graceful fallbacks for missing configuration keys
Example configuration (sft_mlflow.yaml) for SFT training with MLflow
Updated README with installation and usage instructions
Fixed test coverage for MLflow plot logging functionality

The integration follows the existing logger architecture and provides a seamless experience for users who want to use MLflow for experiment tracking alongside or instead of TensorBoard and Weights & Biases.

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

SahilJain314

Thanks for contributing! @terrykong to comment on the feature broadly

README.md

examples/configs/sft_mlflow.yaml

terrykong

thanks @therealnaveenkamal for the contribution! appreciate your help making nemo-rl better!

left some comments below

README.md

examples/configs/sft_mlflow.yaml

nemo_rl/utils/logger.py

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

terrykong

a couple other comments

examples/configs/sft_mlflow.yaml

nemo_rl/utils/logger.py

tests/unit/utils/test_logger.py

therealnaveenkamal · 2025-06-27T23:13:29Z

Sure @terrykong , working on your comments. I'll ping you here once I'm ready. Thanks for your support.

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

therealnaveenkamal · 2025-06-28T03:48:32Z

@terrykong I think I've addressed all your comments. please review and let me know if further edits are required.

examples/configs/grpo_math_1B.yaml

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

therealnaveenkamal · 2025-07-05T22:44:22Z

@terrykong / @SahilJain314 not sure why the workflows have failed. do I have to do anything from my side? the lint & the copyright checks passed when I run the workflows locally.

terrykong · 2025-07-06T04:44:01Z

@therealnaveenkamal I believe the copyright one is orthogonal to your change so you can ignore that (it doesn't gate our CI, but it should CC: @chtruong814)

The linter one fails. can you install the precommit hooks and run it on your branch?

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

therealnaveenkamal · 2025-07-07T01:00:24Z

@terrykong thank you. resolved the lint errors.

terrykong · 2025-07-07T16:06:23Z

thanks @therealnaveenkamal. The CI still says there's some unresolved errors:

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

therealnaveenkamal · 2025-07-14T02:20:18Z

@terrykong sorry my bad, I didn't do a lint check after resolving merge conflicts. i think everything should be good now. thanks.

Signed-off-by: Terry Kong <terryk@nvidia.com>

terrykong · 2025-07-14T18:47:13Z

the linter is failing due to an error that usually indicates the uv.lock needs to be re-locked. i made a PR against your fork to update:

therealnaveenkamal#1

update uv.lock to account for the new dependency

therealnaveenkamal · 2025-07-14T19:08:13Z

Thanks @terrykong, merged.

terrykong · 2025-07-20T05:52:09Z

Hi @therealnaveenkamal . I think the CI is having some issues with your PR (CC @chtruong814 ). I'll open a new PR, resolve the issues, and make sure to credit your contribution

therealnaveenkamal · 2025-07-21T21:02:49Z

Thanks for the support @terrykong !! Looking forward to contributing more.

terrykong · 2025-07-21T21:09:06Z

Appreciate doing most of the heavy lifting @therealnaveenkamal ! I just pushed it over the finish.

I'll close this since #697 has landed

Feat: MLFlow Integration for experiment tracking

890fdc5

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

therealnaveenkamal changed the title ~~Feat: MLFlow Integration for experiment tracking~~ MLFlow Integration for experiment tracking Jun 21, 2025

therealnaveenkamal changed the title ~~MLFlow Integration for experiment tracking~~ feat: MLFlow Integration for experiment tracking Jun 21, 2025

SahilJain314 reviewed Jun 26, 2025

View reviewed changes

README.md Outdated Show resolved Hide resolved

examples/configs/sft_mlflow.yaml Outdated Show resolved Hide resolved

terrykong requested changes Jun 27, 2025

View reviewed changes

therealnaveenkamal added 2 commits June 27, 2025 17:55

code structuring, cleanup and readme update

70e6102

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

code cleanup and structure

4bd8e51

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

github-actions bot added the documentation Improvements or additions to documentation label Jun 27, 2025

terrykong requested changes Jun 27, 2025

View reviewed changes

examples/configs/sft_mlflow.yaml Outdated Show resolved Hide resolved

nemo_rl/utils/logger.py Outdated Show resolved Hide resolved

tests/unit/utils/test_logger.py Outdated Show resolved Hide resolved

therealnaveenkamal added 2 commits June 28, 2025 03:36

code cleanup

ea8fe0f

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

minor readme chaange

47fc68c

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

SahilJain314 requested changes Jun 30, 2025

View reviewed changes

examples/configs/grpo_math_1B.yaml Outdated Show resolved Hide resolved

minor mlflow config updates

2cf0cf0

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

terrykong added the CI:L1 Run doctests, unit tests, and functional tests label Jul 2, 2025

terrykong added CI:docs Run doctest and removed CI:L1 Run doctests, unit tests, and functional tests labels Jul 6, 2025

terrykong temporarily deployed to nemo-ci July 6, 2025 04:40 — with GitHub Actions Inactive

terrykong self-requested a review July 6, 2025 04:43

minor lint fix

9064a45

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

github-actions bot added the CI Relating to CI label Jul 7, 2025

therealnaveenkamal force-pushed the feature/ml-flow-logger branch from e064aec to 9064a45 Compare July 7, 2025 00:12

github-actions bot removed the CI Relating to CI label Jul 7, 2025

Merge branch 'main' into feature/ml-flow-logger

e37ab51

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

terrykong removed the CI:docs Run doctest label Jul 7, 2025

terrykong added the CI:docs Run doctest label Jul 7, 2025

terrykong temporarily deployed to nemo-ci July 7, 2025 16:05 — with GitHub Actions Inactive

therealnaveenkamal added 2 commits July 13, 2025 21:58

fixed lint error

caed641

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

Merge upstream/main: resolve conflicts and integrate MLflow logger

b864de9

update uv.lock to account for the new dependency

90bec34

Signed-off-by: Terry Kong <terryk@nvidia.com>

Merge pull request #1 from NVIDIA-NeMo/tk/uvlock-mlflow

af4bb13

update uv.lock to account for the new dependency

terrykong approved these changes Jul 14, 2025

View reviewed changes

terrykong added CI:L1 Run doctests, unit tests, and functional tests and removed CI:docs Run doctest labels Jul 14, 2025

terrykong temporarily deployed to nemo-ci July 14, 2025 19:50 — with GitHub Actions Inactive

snowmanwwg added the external label Jul 20, 2025

terrykong mentioned this pull request Jul 20, 2025

feat: MLFlow Integration for experiment tracking #697

Merged

terrykong closed this Jul 21, 2025

tpoisonooo mentioned this pull request Jul 23, 2025

feat(logger.py): support swanlab #716

Closed

4 tasks

feat: MLFlow Integration for experiment tracking #534

feat: MLFlow Integration for experiment tracking #534

Uh oh!

Conversation

therealnaveenkamal commented Jun 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

SahilJain314 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

therealnaveenkamal commented Jun 27, 2025

Uh oh!

therealnaveenkamal commented Jun 28, 2025

Uh oh!

Uh oh!

therealnaveenkamal commented Jul 5, 2025

Uh oh!

terrykong commented Jul 6, 2025

Uh oh!

therealnaveenkamal commented Jul 7, 2025

Uh oh!

terrykong commented Jul 7, 2025

Uh oh!

therealnaveenkamal commented Jul 14, 2025

Uh oh!

terrykong commented Jul 14, 2025

Uh oh!

therealnaveenkamal commented Jul 14, 2025

Uh oh!

terrykong commented Jul 20, 2025

Uh oh!

therealnaveenkamal commented Jul 21, 2025

Uh oh!

terrykong commented Jul 21, 2025

Uh oh!

Uh oh!

therealnaveenkamal commented Jun 21, 2025 •

edited

Loading