[Feature] Add modular tracking interface with MLflow backend

[mouad-hpc]

Summary

• Adds TrackingBackend ABC and TrackingManager for pluggable logging backends
• Adds MLflow as a new tracking backend alongside wandb and tensorboard
• Refactors tracking_utils.py to delegate to the manager instead of hardcoded conditionals
• New CLI flags: --use-mlflow, --mlflow-tracking-uri, --mlflow-experiment-name, --mlflow-run-name
• MLflow is an optional dependency (pip install miles[mlflow])

Files changed

• miles/utils/tracking.py (new) - shared interface + adapters + registry
• miles/utils/mlflow_utils.py (new) - MLflow backend implementation
• miles/utils/tracking_utils.py - uses TrackingManager internally, same public API
• miles/utils/arguments.py - MLflow CLI flags
• setup.py - mlflow added as optional extra

Test plan

• Verified e2e functionality on 8xH100 GPU cluster (via slime fork)
• Full training run with --use-mlflow logging metrics end-to-end
• Verify secondary rank attachment via mlflow_run_id
• Verify no regression when mlflow is not installed

[gemini-code-assist]

Summary of Changes

Hello @mouad-hpc, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant architectural improvement to the experiment tracking system by implementing a modular and extensible interface. This change enables seamless integration of various logging backends, with MLflow being the first new addition, alongside existing WandB and Tensorboard support. The refactoring centralizes tracking logic, making the system more flexible and easier to maintain, while also providing new command-line options for MLflow configuration.

Highlights

• Modular Tracking Interface: Introduced a new modular tracking interface using an abstract base class TrackingBackend and a TrackingManager to support pluggable logging backends.
• MLflow Integration: Added MLflow as a new experiment tracking backend, complementing existing WandB and Tensorboard support.
• Refactored Tracking Utilities: Refactored tracking_utils.py to delegate logging operations to the new TrackingManager, simplifying its internal logic and removing hardcoded conditionals.
• New CLI Flags: Implemented new command-line interface flags (--use-mlflow, --mlflow-tracking-uri, --mlflow-experiment-name, --mlflow-run-name, --mlflow-run-id) for configuring MLflow tracking.
• Optional MLflow Dependency: Configured MLflow as an optional dependency, allowing users to install it only if needed via pip install miles[mlflow].

Changelog

• miles/utils/arguments.py
  • Added add_mlflow_arguments function to define new CLI flags for MLflow configuration.
  • Integrated add_mlflow_arguments into the main argument parser.
• miles/utils/mlflow_utils.py
  • Created a new module to encapsulate MLflow-specific initialization, metric logging, and run finalization logic.
  • Implemented helper functions for sanitizing keys and flattening dictionaries for MLflow logging.
• miles/utils/tracking.py
  • Introduced TrackingBackend as an abstract base class defining the interface for all logging backends.
  • Implemented WandbBackend, TensorboardBackend, and MlflowBackend as adapters to existing or new utility functions.
  • Defined BACKEND_REGISTRY to map backend names to their classes and corresponding CLI flags.
  • Created TrackingManager to initialize, log to, and finish multiple active tracking backends.
• miles/utils/tracking_utils.py
  • Replaced direct calls to WandB and Tensorboard utilities with calls to the new TrackingManager instance.
  • Simplified init_tracking, log, and added finish_tracking to delegate to the manager.
• setup.py
  • Added mlflow to the extras_require dictionary, making it an optional installation.

Activity

• No specific activity (comments, reviews, or progress updates) has been provided for this pull request beyond its initial creation.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a well-designed modular tracking interface, which is a significant architectural improvement for extensibility. The implementation of MLflow as a new backend is clean and adheres to best practices for optional dependencies and distributed environments, and the refactoring of tracking_utils.py using TrackingManager effectively simplifies the codebase. However, the new MLflow integration introduces several critical security vulnerabilities, including a Server-Side Request Forgery (SSRF) due to unsanitized tracking URIs, multiple Path Traversal vulnerabilities from unsanitized experiment/run names and IDs, and the logging of sensitive information (W&B API key) to MLflow parameters. These security issues must be addressed by implementing strict validation and sanitization on all user-controlled inputs. Additionally, I have a couple of minor suggestions to further enhance code clarity and robustness.

[yueming-yuan]

a tiny suggestion: could you improve the file structure? maybe put all files under the same subdirectory like utils/tracking

[mouad-hpc]

a tiny suggestion: could you improve the file structure? maybe put all files under the same subdirectory like utils/tracking

Done