Forward mlflow env vars to ray runtime env

skyrl-train/docs/configuration/config.rst

@@ -122,8 +122,8 @@ Logging and Debugging Configuration
     dump_eval_results: true
 
 - ``logger``: Logger to use. Currently, we support ``wandb``, ``mlflow``, and ``console``. ``console`` will simply log metrics to the console. 
-- ``project_name``: Name of the project in WandB.
-- ``run_name``: Name of the run in WandB.
+- ``project_name``: Name of the project in WandB and MLFlow.
+- ``run_name``: Name of the run in WandB and MLFlow.
 - ``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``
 - ``dump_eval_results``: Whether to dump the evaluation results to a file. When ``true``, the full evaluation results will be dumped to a file in the ``export_path`` directory. The evaluation results at global step ``N`` is saved to ``self.cfg.trainer.export_path / "dumped_eval" / global_step_N_eval_results``
 

skyrl-train/skyrl_train/utils/utils.py

@@ -281,12 +281,18 @@ def initialize_ray(cfg: DictConfig):
         env_vars["NCCL_P2P_DISABLE"] = "1"
         env_vars["NCCL_SHM_DISABLE"] = "1"
 
-    # TODO: this can be removed if we standardize on env files.
-    # But it's helpful for a quickstart
     if os.environ.get("WANDB_API_KEY"):
         logger.info("Exporting wandb api key to ray runtime env")
         env_vars["WANDB_API_KEY"] = os.environ["WANDB_API_KEY"]
 
+    if os.environ.get("MLFLOW_TRACKING_URI"):
+        logger.info("Exporting mlflow tracking uri to ray runtime env")
+        env_vars["MLFLOW_TRACKING_URI"] = os.environ["MLFLOW_TRACKING_URI"]
+
+    if os.environ.get("MLFLOW_TRACKING_TOKEN"):
+        logger.info("Exporting mlflow tracking token to ray runtime env")
+        env_vars["MLFLOW_TRACKING_TOKEN"] = os.environ["MLFLOW_TRACKING_TOKEN"]
+
     if os.environ.get("SKYRL_LD_LIBRARY_PATH_EXPORT"):
         # export `LD_LIBRARY_PATH` to ray runtime env.
         # For some reason the `LD_LIBRARY_PATH` is not exported to the worker with .env file.