Add MLflow (Databricks) logging backend

[amazloumi]

Summary

• Adds MLflowBackend, a third _LoggingBackend alongside WandB/TensorBoard, enabled via metrics.enable_mlflow. Logs the same per-step and eval metric dict (train/*, gpu/*, smoothed/*, eval/*, moe/*) to Databricks-hosted MLflow (or any MLflow server).
• Also logs the flattened JobConfig as params + host/slurm_job_id tags + optional system metrics. No artifacts — checkpoints stay on disk.
• Resume: mlflow_run_id is saved in the checkpoint and restored on resume (mirrors wandb_run_id), so a requeue reattaches to the same run — or starts a fresh run if the saved id was deleted (parity with wandb resume="allow").
• Credentials are env-only (DATABRICKS_HOST + DATABRICKS_TOKEN/DATABRICKS_API_TOKEN); experiment resolves mlflow_experiment → $MLFLOW_EXPERIMENT → auto.
• Optional dep group mlflow (mlflow-skinny + databricks-sdk) keeps the base install lean and avoids a pandas downgrade; adds a check_env --requires mlflow preflight tag and docs/metrics-and-profiling/mlflow.md.
• Default off; degrades to a warning if the dep/creds are absent. In-training eval logs automatically via log_eval; the standalone eval pipeline is unchanged (no backend — same as WandB/TB).

Testing

• uv run ruff check kempnerforge/ tests/ passes
• uv run ruff format --check kempnerforge/ tests/ scripts/ passes
• uv run pyright kempnerforge/ passes (0 errors)
• uv run pytest tests/unit/ --timeout=60 passes (1658 passed, 3 skipped)
• Distributed — N/A (no distributed code changed)
• E2E — not run (no GPU). The scripts/train.py change is a 3-line additive mirror of the existing wandb_run_id checkpoint handling. Separately verified live: real logging to Databricks MLflow, plus resume / delete→new-run / system metrics via an end-to-end run against real MLflow.

Closes #158

[Copilot]

Pull request overview

Adds an optional MLflow (Databricks-compatible) metrics logging backend to the training metrics system, alongside the existing WandB and TensorBoard integrations, including config/doc updates and checkpoint-based run resume support.

Changes:

• Introduces MLflowBackend with lazy initialization, run-id writeback for resume, and flattened JobConfig param logging.
• Adds [metrics] configuration for MLflow + scripts/train.py checkpoint persistence/restoration of mlflow_run_id.
• Adds check_env --requires mlflow support, docs, tests, and optional dependency group (mlflow-skinny, databricks-sdk, etc.).

Reviewed changes

Copilot reviewed 13 out of 15 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
kempnerforge/metrics/tracker.py	Adds MLflow backend, experiment resolution, and config flattening; wires backend selection into MetricsTracker.
kempnerforge/config/metrics.py	Adds MLflow config fields + Databricks experiment absolute-path validation for config-provided experiment.
scripts/train.py	Persists/restores mlflow_run_id in checkpoint metadata to support resume/requeue.
scripts/check_env.py	Adds mlflow preflight tag for Databricks credential presence (and optional SDK probe).
tests/unit/test_observability.py	Adds unit tests for MLflow backend behavior, config flattening, and Databricks experiment-path validation.
pyproject.toml	Adds optional mlflow dependency group.
uv.lock	Locks new optional dependencies (mlflow-skinny, databricks-sdk, and transitive deps) and updates protobuf.
docs/metrics-and-profiling/mlflow.md	New documentation page describing MLflow backend usage, config, and failure behavior.
docs/metrics-and-profiling/metrics-tracker.md	Updates backend construction and close semantics to include MLflow.
docs/metrics-and-profiling/index.md	Updates module overview/config examples to include MLflow.
docs/configuration/config-sections.md	Documents new [metrics] MLflow fields.
README.md	Mentions MLflow as a supported backend and updates diagram/module listing.
configs/train/debug.toml	Adds commented MLflow example config hints.
CHANGELOG.md	Records addition of MLflow backend and related changes.
.gitignore	Ignores local MLflow SQLite store + mlruns/ directory.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

❌ Patch coverage is 94.96855% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
kempnerforge/metrics/tracker.py	94.66%	5 Missing and 3 partials ⚠️

Files with missing lines	Coverage Δ
kempnerforge/config/metrics.py	100.00% <100.00%> (ø)
kempnerforge/metrics/tracker.py	91.80% <94.66%> (+2.99%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[camilobrownpinilla]

One potential security concern, but that is already tracked as a separate issue. Bringing it up here for awareness, but the PR looks good otherwise.

[mmshad]

Ready to merge. Only minor issues:

  1. # pyright: reportMissingImports=false is applied file-scope in tracker.py. It silences the optional mlflow imports, but also any genuinely missing import elsewhere in that file. Scoping it to the
  lazy-import lines would be tighter.

  2. _log_run_metadata batches params 100 at a time, but the try/except wraps the whole loop: if one batch is rejected (e.g. an over-long nested key), the remaining param batches are abandoned.
  Non-fatal, just missing params, and it warns.

  3. The dedup watermark hardcodes "train/loss". Correct today, but a future rename of that key would silently break resume-dedup with no test catching it (every test injects history directly).