-
Notifications
You must be signed in to change notification settings - Fork 1
[feat] MLflow Integration #14
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
Co-authored-by: Copilot <[email protected]>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
Adds MLflow experiment tracking to ReAgentAI, updates the UI and main app to log metrics and parameters, and provides Docker and documentation support for MLflow.
- Introduces
MLflowTrackerfor run management and logging - Updates
main.pyand UI handlers to start runs and log to MLflow when enabled - Adds Dockerfile, docker-compose setup, and documentation for running MLflow alongside the app
Reviewed Changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_reagentai.py | Adds tests for MLflowTracker initialization, lifecycle, and error paths |
| src/reagentai/common/mlflow_tracking.py | Implements the MLflowTracker class with start/end run and logging methods |
| src/reagentai/ui/app.py | Extends UI callbacks to accept an optional mlflow_tracker and log data |
| src/reagentai/main.py | Initializes and uses MLflowTracker, logging system info and ending runs |
| pyproject.toml | Adds mlflow to dependencies and pytest configuration |
| docker-compose.yml | Defines services for running ReAgentAI and MLflow together |
| README.md | Updates setup and run instructions to include MLflow usage |
| Dockerfile.mlflow | Provides a container image for the MLflow server |
Comments suppressed due to low confidence (4)
README.md:58
- [nitpick] The command
uv run mlflow serverappears incorrect; it should likely bemlflow server(or use the correct CLI tool) without theuv runprefix.
uv run mlflow server
src/reagentai/ui/app.py:96
- The variable
user_queryis not defined inhandle_bot_response. You should either includeuser_queryin the function signature or extract it from inputs before using it.
f"query_{len(chat_history)}": user_query[:100] # Truncate long queries
src/reagentai/main.py:23
loggeris not defined in this module; replacelogger.infowithlogging.infoor initialize a module-level logger vialogger = logging.getLogger(__name__).
logger.info(f"MLflow tracking {'enabled' if run_id else 'disabled'}")
src/reagentai/main.py:59
loggeris not defined here either; uselogging.infoor definelogger = logging.getLogger(__name__)at the top of the file.
logger.info("MLflow tracking session ended")
| [tool.pytest.ini_options] | ||
| testpaths = ["tests"] | ||
| python_files = "test_*.py" | ||
|
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe add this? Because it doesn't work without that for me (the reagent module is not found when running tests).
pythonpath = [
"."
]
| networks: | ||
| - mlflow-network | ||
| networks: | ||
| mlflow-network: No newline at end of file |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nitpick: missing new line at the end
|
|
||
| # Log metrics to MLflow | ||
| if mlflow_tracker and mlflow_tracker.mlflow_enabled: | ||
| mlflow_tracker.log_metrics( | ||
| {"token_usage": token_used, "conversation_length": len(chat_history)} | ||
| ) | ||
|
|
||
| # Log user query as param for tracking purposes | ||
| mlflow_tracker.log_params( | ||
| { | ||
| f"query_{len(chat_history)}": user_query[:100] # Truncate long queries | ||
| } | ||
| ) | ||
|
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I wouldn't pass the mlflow instance to the Gradio frontend, I don't think we should keep logic like that there.
Maybe we should pass the tracker to the Agent class itself.
| # Initialize MLflow tracking | ||
| tracker = MLflowTracker(experiment_name="reagentai_experiments") | ||
|
|
||
| # Start a new run for this application session | ||
| run_id = tracker.start_run(run_name="reagentai_session") | ||
| logger.info(f"MLflow tracking {'enabled' if run_id else 'disabled'}") | ||
|
|
||
| # Log system information and configuration parameters | ||
| if tracker.mlflow_enabled: | ||
| import platform | ||
| import sys | ||
|
|
||
| # Log system info as tags | ||
| tracker.set_tags( | ||
| { | ||
| "python_version": sys.version, | ||
| "platform": platform.platform(), | ||
| "application": "ReagentAI", | ||
| } | ||
| ) | ||
|
|
||
| # Log configuration parameters | ||
| tracker.log_params( | ||
| { | ||
| "log_to_file": os.environ.get("LOG_TO_FILE", "True"), | ||
| "app_version": "0.1.0", # Could be pulled from a version file | ||
| } | ||
| ) | ||
|
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe add a function like setup_mlflow for this?
Summary
This PR adds MLflow integration to the ReAgentAI project, enabling experiment tracking and visualization. The integration includes:
MLflowTrackerclass for tracking experiments, metrics, and artifactsChanges
Features