mlflow + trackio logging

[humaira-rf]

Changes

• Added mlflow logging to evals
• Added trackio logging to both fit and evals
• Deprecated RF_TRACKING_BACKEND environment variable for backends
• Environment variable RF_MLFLOW_ENABLED to enable MLFLOW tracking (default to true for non-Colab environments)
• Environment variable RF_TENSORBOARD_ENABLED to enable TENSORBOARD tracking (default to true for Colab environments)
• Environment variable RF_TRACKIO_ENABLED to enable TRACKIO tracking (default to false for all environments)
• --tracking-backend flag for rapidfireai start changed to --tracking-backends to now allow multiple flags with one of the following: mlflow, tensorboard, trackio
• Remove more pinning of dependent modules
• Rename mlflow_run_id column in database to metric_run_id to be more generic for metrics logger
• Rename mlflow_experiment_id column in databasse to metric_experment_id to be more generic for metrics logger
• Consolidated mlflow_manager, tensorboard_manager, and trackio_manager for both evals and fit
• New default metric_rfmetric_manager as Metrics manager that accepts one or more Metrics loggers

---

Note

Introduces a unified metrics layer and migrates the codebase off MLflow-only assumptions.

• Add RFMetricLogger with backends: MLflow, TensorBoard, TrackIO; implement per-backend managers and default selection via RF_MLFLOW_ENABLED, RF_TENSORBOARD_ENABLED, RF_TRACKIO_ENABLED
• Replace MLflow-specific fields/logic across evals/fit with generic metrics: DB schema migrations (mlflow_* → metric_*), controller/worker/experiment now create/log/end runs via metric_manager
• CLI: --tracking-backends (multi-select) replaces --tracking-backend; sets new env flags; add --force for non-interactive shutdown
• Callbacks: rename to MetricLoggingCallback; Generation metrics use torch.amp.autocast("cuda"); trainer and configs updated to pass metric_run_id
• Remove legacy MLflow-only helpers; add new metric managers (mlflow, tensorboard, trackio) and abstraction types
• Update README (directory layout, env vars), start script and doctor output to new flags; loosen/select deps (add trackio, optionalize mlflow, set RF_TENSORBOARD_LOG_DIR default)
• Evals/fit DB schema files and runtime add migrations; tests adjusted to new interfaces

^{Written by Cursor Bugbot for commit eb6153f. This will update automatically on new commits. Configure here.}

• Tested on GCP and Colab for both fit and evals

[david-rfai]

See comments, and also look into other ways of combining evals and fit functions/classes. Would also be nice if the database could abstract the metrics logger, so there does not need to be dedicated columns for different trackers.

[cursor]

This PR is being reviewed by Cursor Bugbot

Details

Your team is on the Bugbot Free tier. On this plan, Bugbot will review limited PRs each billing cycle for each member of your team.

To receive Bugbot reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial.

resolved

[cursor]

This PR is being reviewed by Cursor Bugbot

Details

Your team is on the Bugbot Free tier. On this plan, Bugbot will review limited PRs each billing cycle for each member of your team.

To receive Bugbot reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial.

[cursor]

This PR is being reviewed by Cursor Bugbot

Details

Your team is on the Bugbot Free tier. On this plan, Bugbot will review limited PRs each billing cycle for each member of your team.

To receive Bugbot reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial.

[arun-rfai]

Nice work!