MLflow/WandB/Comet/Trackio env vars not set on Ray Train worker (use_ray: true)

Bug

When use_ray: true, tracking integrations (MLflow, WandB, Comet, Trackio) silently fail to log metrics because their environment variables are never set on the Ray Train worker process.

Root cause

load_cfg() in cli/config.py (lines 324-327) calls:

setup_wandb_env_vars(cfg)
setup_mlflow_env_vars(cfg)
setup_comet_env_vars(cfg)
setup_trackio_env_vars(cfg)

These functions convert YAML config fields (mlflow_tracking_uri, mlflow_experiment_name, wandb_api_key, etc.) into environment variables that the HuggingFace callbacks (MLflowCallback, WandbCallback, etc.) read via os.getenv().

However, load_cfg() runs on the driver process. When use_ray: true, the actual training runs on a separate Ray Train worker process via ray_train_func() in cli/train.py. The worker reconstructs the config from the serialized dict but never calls these setup functions. The HF callbacks read env vars that don't exist, so tracking silently fails.

Without use_ray (direct accelerate/torchrun path), load_cfg and training run in the same process, so it works.

Impact

• mlflow_tracking_uri, mlflow_experiment_name from YAML are ignored → MLflow logs to "Default" experiment or doesn't log at all
• Same for wandb_*, comet_*, trackio_* config fields
• No error is raised — tracking just silently doesn't work

Prior context

This was flagged during review of PR #3130 by @NanoCode012:

https://github.com/axolotl-ai-cloud/axolotl/pull/3130#discussion_r2329564517

Do we need any of the others too?

python复制

setup_wandb_env_vars(cfg)
setup_mlflow_env_vars(cfg)
setup_comet_env_vars(cfg)

The question was not addressed and PR #3130 was merged without adding these calls.

Fix

Add the four setup calls in ray_train_func() in cli/train.py, mirroring what load_cfg() does on the driver. See PR: https://github.com/axolotl-ai-cloud/axolotl/pull/3909