Skip to content

Conversation

@tdadela
Copy link
Collaborator

@tdadela tdadela commented Jun 8, 2025

Summary

This PR adds MLflow integration to the ReAgentAI project, enabling experiment tracking and visualization. The integration includes:

  1. A new MLflowTracker class for tracking experiments, metrics, and artifacts
  2. Docker configuration for running MLflow alongside ReAgentAI
  3. Updated documentation and setup instructions

Changes

Features

  • Added MLflow integration for experiment tracking
  • Created Dockerfile.mlflow for running MLflow server in Docker
  • Added docker-compose.yml for running both ReAgentAI and MLflow together
  • Updated pyproject.toml to include MLflow dependency

@tdadela tdadela requested a review from Copilot June 8, 2025 10:19

This comment was marked as outdated.

tdadela and others added 2 commits June 8, 2025 12:23
@tdadela tdadela changed the title MLflow Integration [feat] MLflow Integration Jun 8, 2025
@tdadela tdadela mentioned this pull request Jun 8, 2025
@tdadela tdadela requested a review from Copilot June 8, 2025 10:33

This comment was marked as outdated.

@tdadela tdadela requested review from Kowalski1024 and janumiko and removed request for Kowalski1024 and janumiko June 8, 2025 10:34
@tdadela tdadela requested a review from Copilot June 8, 2025 11:36

This comment was marked as outdated.

@tdadela tdadela requested review from Copilot and janumiko June 8, 2025 13:32
Copy link
Contributor

Copilot AI left a 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 MLflowTracker for run management and logging
  • Updates main.py and 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 server appears incorrect; it should likely be mlflow server (or use the correct CLI tool) without the uv run prefix.
uv run mlflow server

src/reagentai/ui/app.py:96

  • The variable user_query is not defined in handle_bot_response. You should either include user_query in 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

  • logger is not defined in this module; replace logger.info with logging.info or initialize a module-level logger via logger = logging.getLogger(__name__).
logger.info(f"MLflow tracking {'enabled' if run_id else 'disabled'}")

src/reagentai/main.py:59

  • logger is not defined here either; use logging.info or define logger = logging.getLogger(__name__) at the top of the file.
logger.info("MLflow tracking session ended")

Comment on lines +20 to +23
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = "test_*.py"

Copy link
Owner

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
Copy link
Owner

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

Comment on lines +86 to +99

# 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
}
)

Copy link
Owner

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.

Comment on lines +18 to +46
# 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
}
)

Copy link
Owner

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?

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.

3 participants