Skip to content

Commit 19b8ca5

Browse files
authored
Fix MLFlow logging (#121)
This is a small change to make the MLFlow integration work. Currently this fails with a Pandas error when trying to flatten an Omega dict; we need to convert to a regular Python dictionary. Can confirm this works on our MLFlow setup: <img width="1406" height="683" alt="image" src="https://github.com/user-attachments/assets/fcee526a-815e-4f08-bf25-d2709779ced7" />
1 parent 321b8fa commit 19b8ca5

File tree

2 files changed

+5
-1
lines changed

2 files changed

+5
-1
lines changed

skyrl-train/docs/configuration/config.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ Logging and Debugging Configuration
121121
dump_data_batch: false
122122
dump_eval_results: true
123123
124-
- ``logger``: Logger to use. Currently, we support ``wandb`` and ``console``. ``console`` will simply log metrics to the console.
124+
- ``logger``: Logger to use. Currently, we support ``wandb``, ``mlflow``, and ``console``. ``console`` will simply log metrics to the console.
125125
- ``project_name``: Name of the project in WandB.
126126
- ``run_name``: Name of the run in WandB.
127127
- ``dump_data_batch``: Whether to dump the data batch to a file. This is useful for debugging. When ``true``, the data batch will be dumped to a file in the ``export_path`` directory. The training batch at global step ``N`` is saved to ``self.cfg.trainer.export_path / "dumped_data" / global_step_N_training_input``

skyrl-train/skyrl_train/utils/tracking.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from pathlib import Path
2222
from typing import Any, Dict, List, Union
2323
from loguru import logger
24+
from omegaconf import DictConfig, OmegaConf
2425
import pprint
2526

2627

@@ -184,6 +185,9 @@ def _compute_mlflow_params_from_objects(params) -> Dict[str, Any]:
184185
if params is None:
185186
return {}
186187

188+
if isinstance(params, DictConfig):
189+
params = OmegaConf.to_container(params, resolve=True)
190+
187191
return _flatten_dict(_transform_params_to_json_serializable(params, convert_list_to_dict=True), sep="/")
188192

189193

0 commit comments

Comments
 (0)