refactor: match functional dir style for wandb dir

alex-petrenko · Nov 18, 2024 · e5760d2 · e5760d2
1 parent 8f3bb78
commit e5760d2
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 4 deletions.
diff --git a/sample_factory/cfg/cfg.py b/sample_factory/cfg/cfg.py
@@ -746,9 +746,8 @@ def add_wandb_args(p: ArgumentParser):
     )
     p.add_argument(
         "--wandb_dir",
-        default="wandb",
+        default=join(os.getcwd(), "wandb"),
         type=str,
-        nargs="*",
         help="Logging Directory for WandB",
     )
 

diff --git a/sample_factory/utils/utils.py b/sample_factory/utils/utils.py
@@ -402,6 +402,10 @@ def experiments_dir(cfg, mkdir=True) -> str:
     return maybe_ensure_dir_exists(cfg.train_dir, mkdir)
 
 
+def wandb_dir(cfg, mkdir=True) -> str:
+    return maybe_ensure_dir_exists(cfg.wandb_dir, mkdir)
+
+
 def experiment_dir(cfg, mkdir=True) -> str:
     experiment = cfg.experiment
     experiments_root = experiments_dir(cfg, mkdir)

diff --git a/sample_factory/utils/wandb_utils.py b/sample_factory/utils/wandb_utils.py
@@ -1,6 +1,6 @@
 from datetime import datetime
 
-from sample_factory.utils.utils import log, retry
+from sample_factory.utils.utils import log, retry, wandb_dir
 
 
 def init_wandb(cfg):
@@ -44,7 +44,7 @@ def init_wandb_func():
             tags=cfg.wandb_tags,
             resume="allow",
             settings=wandb.Settings(start_method="fork"),
-            dir=cfg.wandb_dir,
+            dir=wandb_dir(cfg, True),
         )
 
     log.debug("Initializing WandB...")