allow additional settings to WandBWriter

GitOrigin-RevId: ba40089

Author: Yunlong Xu

axlearn/common/summary_writer.py

@@ -402,6 +402,10 @@ class Config(BaseWriter.Config):
         # If True, convert any 2D Tensors to wandb.Image before passing to wandb.log.
         convert_2d_to_image: bool = False
 
+        # Dictionary kwargs to pass to wandb.Settings().
+        # If None, no additional custom settings will be configured.
+        wandb_settings_kwargs: Optional[dict[str, Any]] = None
+
     @classmethod
     def default_config(cls: Config) -> Config:
         cfg = super().default_config()
@@ -444,6 +448,11 @@ def _initialize_run(self):
             with fs.open(wandb_file, "w") as f:  # pytype: disable=module-attr
                 f.write(exp_id)
 
+        # Build custom settings from configured kwargs
+        wandb_settings = (
+            wandb.Settings(**cfg.wandb_settings_kwargs) if cfg.wandb_settings_kwargs else None
+        )
+
         wandb.init(
             id=exp_id,
             name=cfg.exp_name,
@@ -456,6 +465,7 @@ def _initialize_run(self):
             resume=cfg.resume,
             dir=cfg.dir,
             group=cfg.group,
+            settings=wandb_settings,
         )
 
     @staticmethod

axlearn/common/summary_writer_test.py

@@ -321,6 +321,28 @@ def test_resume(self):
             finally:
                 wandb.finish()
 
+    @pytest.mark.skipif(wandb is None, reason="wandb package not installed.")
+    @pytest.mark.skipif("WANDB_API_KEY" not in os.environ, reason="wandb api key not found.")
+    def test_wandb_settings_kwargs(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            try:
+                _: WandBWriter = (
+                    WandBWriter.default_config()
+                    .set(
+                        name="test",
+                        exp_name="wandb-testSettings",
+                        dir=tempdir,
+                        mode="offline",
+                        wandb_settings_kwargs={"silent": True, "disable_code": True},
+                    )
+                    .instantiate(parent=None)
+                )
+                # Verify the settings were applied
+                self.assertTrue(wandb.run.settings.silent)
+                self.assertTrue(wandb.run.settings.disable_code)
+            finally:
+                wandb.finish()
+
 
 if __name__ == "__main__":
     absltest.main()