初始化项目，由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
ModelHub XC
228
slime/utils/wandb_utils.py
Normal file
228
slime/utils/wandb_utils.py
Normal file
@@ -0,0 +1,228 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import logging
|
||||
import os
|
||||
from copy import deepcopy
|
||||
|
||||
import wandb
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _is_offline_mode(args) -> bool:
|
||||
"""Detect whether W&B should run in offline mode.
|
||||
|
||||
Priority order:
|
||||
1) args.wandb_mode if provided
|
||||
2) WANDB_MODE environment variable
|
||||
"""
|
||||
if args.wandb_mode:
|
||||
return args.wandb_mode == "offline"
|
||||
return os.environ.get("WANDB_MODE") == "offline"
|
||||
|
||||
|
||||
def init_wandb_primary(args):
|
||||
if not args.use_wandb:
|
||||
args.wandb_run_id = None
|
||||
return
|
||||
|
||||
# Set W&B mode if specified (overrides WANDB_MODE env var)
|
||||
if args.wandb_mode:
|
||||
os.environ["WANDB_MODE"] = args.wandb_mode
|
||||
if args.wandb_mode == "offline":
|
||||
logger.info("W&B offline mode enabled. Data will be saved locally.")
|
||||
elif args.wandb_mode == "disabled":
|
||||
logger.info("W&B disabled mode enabled. No data will be logged.")
|
||||
elif args.wandb_mode == "online":
|
||||
logger.info("W&B online mode enabled. Data will be uploaded to cloud.")
|
||||
|
||||
offline = _is_offline_mode(args)
|
||||
|
||||
# Only perform explicit login when NOT offline
|
||||
if (not offline) and args.wandb_key is not None:
|
||||
wandb.login(key=args.wandb_key, host=args.wandb_host)
|
||||
|
||||
# Check if we should resume a previous run
|
||||
# Priority: 1) wandb_resume_run_id from args, 2) wandb_run_id from args, 3) wandb_run_id from checkpoint file
|
||||
resume_run_id = getattr(args, "wandb_resume_run_id", None) or getattr(args, "wandb_run_id", None)
|
||||
if not resume_run_id:
|
||||
resume_run_id = _load_wandb_run_id_from_checkpoint(args)
|
||||
|
||||
if resume_run_id:
|
||||
# Resume an existing run
|
||||
logger.info(f"Resuming W&B run with id: {resume_run_id}")
|
||||
init_kwargs = {
|
||||
"id": resume_run_id,
|
||||
"entity": args.wandb_team,
|
||||
"project": args.wandb_project,
|
||||
"resume": "must", # Fail if run doesn't exist
|
||||
"config": _compute_config_for_logging(args),
|
||||
}
|
||||
|
||||
# Configure settings based on offline/online mode
|
||||
if offline:
|
||||
init_kwargs["settings"] = wandb.Settings(mode="offline")
|
||||
else:
|
||||
init_kwargs["settings"] = wandb.Settings(mode="shared", x_primary=True)
|
||||
else:
|
||||
# Create a new run
|
||||
# add random 6 length string with characters
|
||||
if args.wandb_random_suffix:
|
||||
suffix = "_" + wandb.util.generate_id()
|
||||
max_base_len = 128 - len(suffix)
|
||||
group = args.wandb_group[:max_base_len] + suffix
|
||||
run_name = f"{group}-RANK_{args.rank}"
|
||||
else:
|
||||
group = args.wandb_group
|
||||
run_name = args.wandb_group
|
||||
|
||||
# Prepare wandb init parameters
|
||||
init_kwargs = {
|
||||
"entity": args.wandb_team,
|
||||
"project": args.wandb_project,
|
||||
"group": group,
|
||||
"name": run_name,
|
||||
"config": _compute_config_for_logging(args),
|
||||
}
|
||||
|
||||
# Configure settings based on offline/online mode
|
||||
if offline:
|
||||
init_kwargs["settings"] = wandb.Settings(mode="offline")
|
||||
else:
|
||||
init_kwargs["settings"] = wandb.Settings(mode="shared", x_primary=True)
|
||||
|
||||
# Add custom directory if specified
|
||||
if args.wandb_dir:
|
||||
# Ensure directory exists to avoid backend crashes
|
||||
os.makedirs(args.wandb_dir, exist_ok=True)
|
||||
init_kwargs["dir"] = args.wandb_dir
|
||||
logger.info(f"W&B logs will be stored in: {args.wandb_dir}")
|
||||
|
||||
wandb.init(**init_kwargs)
|
||||
|
||||
_init_wandb_common()
|
||||
|
||||
args.wandb_run_id = wandb.run.id
|
||||
_save_wandb_run_id_to_checkpoint(args)
|
||||
|
||||
if resume_run_id:
|
||||
logger.info(f"Successfully resumed W&B run: {wandb.run.url}")
|
||||
|
||||
|
||||
def _load_wandb_run_id_from_checkpoint(args):
|
||||
load_dir = getattr(args, "load", None)
|
||||
if not load_dir:
|
||||
return None
|
||||
path = os.path.join(load_dir, "wandb_run_id.txt")
|
||||
if os.path.exists(path):
|
||||
with open(path, "r") as f:
|
||||
run_id = f.read().strip()
|
||||
if run_id:
|
||||
logger.info(f"Loaded wandb_run_id from {path}: {run_id}")
|
||||
return run_id
|
||||
return None
|
||||
|
||||
|
||||
def _save_wandb_run_id_to_checkpoint(args):
|
||||
save_dir = getattr(args, "save", None)
|
||||
if not save_dir:
|
||||
return
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
path = os.path.join(save_dir, "wandb_run_id.txt")
|
||||
with open(path, "w") as f:
|
||||
f.write(args.wandb_run_id)
|
||||
logger.info(f"Saved wandb_run_id to {path}")
|
||||
|
||||
|
||||
def _compute_config_for_logging(args):
|
||||
output = deepcopy(args.__dict__)
|
||||
|
||||
whitelist_env_vars = [
|
||||
"SLURM_JOB_ID",
|
||||
# We may insert more default values here, and may also allow users to configure a whitelist
|
||||
]
|
||||
output["env_vars"] = {k: v for k, v in os.environ.items() if k in whitelist_env_vars}
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# https://docs.wandb.ai/guides/track/log/distributed-training/#track-all-processes-to-a-single-run
|
||||
def init_wandb_secondary(args, router_addr=None):
|
||||
wandb_run_id = getattr(args, "wandb_run_id", None)
|
||||
if wandb_run_id is None:
|
||||
return
|
||||
|
||||
# Set W&B mode if specified (same as primary)
|
||||
if args.wandb_mode:
|
||||
os.environ["WANDB_MODE"] = args.wandb_mode
|
||||
|
||||
offline = _is_offline_mode(args)
|
||||
|
||||
if (not offline) and args.wandb_key is not None:
|
||||
wandb.login(key=args.wandb_key, host=args.wandb_host)
|
||||
|
||||
# Configure settings based on offline/online mode
|
||||
if offline:
|
||||
settings_kwargs = dict(mode="offline")
|
||||
else:
|
||||
settings_kwargs = dict(
|
||||
mode="shared",
|
||||
x_primary=False,
|
||||
x_update_finish_state=False,
|
||||
)
|
||||
|
||||
if args.sglang_enable_metrics and router_addr is not None:
|
||||
logger.info(f"Forward SGLang metrics at {router_addr} to WandB.")
|
||||
settings_kwargs |= dict(
|
||||
x_stats_open_metrics_endpoints={
|
||||
"sgl_engine": f"{router_addr}/engine_metrics",
|
||||
},
|
||||
x_stats_open_metrics_filters={
|
||||
"sgl_engine.*": {},
|
||||
},
|
||||
)
|
||||
|
||||
init_kwargs = {
|
||||
"id": wandb_run_id,
|
||||
"entity": args.wandb_team,
|
||||
"project": args.wandb_project,
|
||||
"config": args.__dict__,
|
||||
"resume": "allow",
|
||||
"reinit": True,
|
||||
"settings": wandb.Settings(**settings_kwargs),
|
||||
}
|
||||
|
||||
# Add custom directory if specified
|
||||
if args.wandb_dir:
|
||||
os.makedirs(args.wandb_dir, exist_ok=True)
|
||||
init_kwargs["dir"] = args.wandb_dir
|
||||
|
||||
wandb.init(**init_kwargs)
|
||||
|
||||
_init_wandb_common()
|
||||
|
||||
|
||||
def _init_wandb_common():
|
||||
wandb.define_metric("train/step")
|
||||
wandb.define_metric("train/*", step_metric="train/step")
|
||||
wandb.define_metric("rollout/step")
|
||||
wandb.define_metric("rollout/*", step_metric="rollout/step")
|
||||
wandb.define_metric("multi_turn/*", step_metric="rollout/step")
|
||||
wandb.define_metric("passrate/*", step_metric="rollout/step")
|
||||
wandb.define_metric("eval/step")
|
||||
wandb.define_metric("eval/*", step_metric="eval/step")
|
||||
wandb.define_metric("perf/*", step_metric="rollout/step")
|
||||
|
||||
|
||||
def get_wandb_offline_dir(args):
|
||||
"""Get the directory where offline W&B data is stored."""
|
||||
if _is_offline_mode(args):
|
||||
if args and hasattr(args, "wandb_dir") and args.wandb_dir:
|
||||
# Use custom directory if specified
|
||||
return args.wandb_dir
|
||||
else:
|
||||
# Default offline directory is ~/wandb/offline-run-<timestamp>
|
||||
# This will be created automatically by wandb
|
||||
return os.path.expanduser("~/wandb")
|
||||
return None
|
||||
Reference in New Issue
Block a user