[Metrics] Allow getting wandb project name via WANDB_PROJECT (#888)

yzhangcs · web-flow · commit 6cb13c7a97ee · 2025-02-25T22:28:51.000-08:00
This commit 
1. Move the `tensorboard` imports into class so that it will not raise
errors if `tensorboard` not installed.
2. Allow users to obtain `wandb` project name from the env variable
`WANDB_PROJECT`, which is defined in their docs.
diff --git a/torchtitan/tools/metrics.py b/torchtitan/tools/metrics.py
@@ -137,7 +137,7 @@ def __init__(self, log_dir: str, tag: Optional[str] = None):
         os.makedirs(log_dir, exist_ok=True)
 
         self.wandb.init(
-            project="torchtitan",
+            project=os.getenv("WANDB_PROJECT", "torchtitan"),
             dir=log_dir,
         )
         logger.info("WandB logging enabled")

Original file line number	Diff line number	Diff line change
`@@ -137,7 +137,7 @@ def __init__(self, log_dir: str, tag: Optional[str] = None):`
`137`	`137`	`os.makedirs(log_dir, exist_ok=True)`
`138`	`138`
`139`	`139`	`self.wandb.init(`
`140`		`- project="torchtitan",`
	`140`	`+ project=os.getenv("WANDB_PROJECT", "torchtitan"),`
`141`	`141`	`dir=log_dir,`
`142`	`142`	`)`
`143`	`143`	`logger.info("WandB logging enabled")`