feat: add wandb offline mode and custom wandb directory support

Author: bingyang-lei

specforge/args.py

@@ -11,6 +11,8 @@ class TrackerArgs:
     wandb_project: str = None
     wandb_name: str = None
     wandb_key: str = None
+    wandb_offline: bool = False
+    wandb_dir: str = None
     swanlab_project: str = None
     swanlab_name: str = None
     swanlab_key: str = None
@@ -33,6 +35,17 @@ def add_args(parser: argparse.ArgumentParser) -> None:
         parser.add_argument("--wandb-project", type=str, default=None)
         parser.add_argument("--wandb-name", type=str, default=None)
         parser.add_argument("--wandb-key", type=str, default=None, help="W&B API key.")
+        parser.add_argument(
+            "--wandb-offline",
+            action="store_true",
+            help="Enable W&B offline mode and store logs locally.",
+        )
+        parser.add_argument(
+            "--wandb-dir",
+            type=str,
+            default=None,
+            help="Directory to store W&B files. Defaults to './wandb' under the project root when using W&B.",
+        )
         # swanlab-specific args
         parser.add_argument(
             "--swanlab-project",

specforge/tracker.py

@@ -91,13 +91,24 @@ def close(self):
 class WandbTracker(Tracker):
     """Tracks experiments using Weights & Biases."""
 
+    @staticmethod
+    def _default_wandb_dir() -> str:
+        # specforge/tracker.py -> project root is one level up
+        return os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "wandb"))
+
     @classmethod
     def validate_args(cls, parser, args):
         if wandb is None:
             parser.error(
                 "To use --report-to wandb, you must install wandb: 'pip install wandb'"
             )
 
+        if args.wandb_dir is None:
+            args.wandb_dir = cls._default_wandb_dir()
+
+        if args.wandb_offline:
+            return
+
         if args.wandb_key is not None:
             return
 
@@ -128,10 +139,24 @@ def validate_args(cls, parser, args):
     def __init__(self, args, output_dir: str):
         super().__init__(args, output_dir)
         if self.rank == 0:
-            wandb.login(key=args.wandb_key)
-            wandb.init(
-                project=args.wandb_project, name=args.wandb_name, config=vars(args)
-            )
+            if args.wandb_dir is None:
+                args.wandb_dir = self._default_wandb_dir()
+            os.makedirs(args.wandb_dir, exist_ok=True)
+            if args.wandb_offline:
+                os.environ["WANDB_MODE"] = "offline"
+            os.environ["WANDB_DIR"] = args.wandb_dir
+
+            if not args.wandb_offline:
+                wandb.login(key=args.wandb_key)
+            init_kwargs = {
+                "project": args.wandb_project,
+                "name": args.wandb_name,
+                "config": vars(args),
+                "dir": args.wandb_dir,
+            }
+            if args.wandb_offline:
+                init_kwargs["mode"] = "offline"
+            wandb.init(**init_kwargs)
             self.is_initialized = True
 
     def log(self, log_dict: Dict[str, Any], step: Optional[int] = None):