fixing wandb cache clean paths (#202)

Author: dylanuys

VERSION

@@ -1 +1 @@
-3.0.2
+3.0.3

bitmind/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.0.2"
+__version__ = "3.0.3"
 
 version_split = __version__.split(".")
 __spec_version__ = (

bitmind/wandb_utils.py

@@ -267,10 +267,14 @@ def clean_wandb_cache(wandb_dir, hours=1):
         wandb_dir: Directory containing wandb run files
         hours: Number of hours to keep runs for (default: 1)
     """
+    if not wandb_dir.endswith("wandb"):
+        wandb_dir = os.path.join(wandb_dir, "wandb")
+
     if not os.path.exists(wandb_dir):
         bt.logging.warning(f"W&B directory not found: {wandb_dir}")
         return
 
+    bt.logging.info(f"Attempting to clean wandb cache at {wandb_dir}")
     run_dirs = [
         d for d in glob.glob(os.path.join(wandb_dir, "run-*")) if os.path.isdir(d)
     ]

neurons/generator.py

@@ -36,7 +36,7 @@
 import bittensor as bt
 from bitmind.config import add_args, add_data_generator_args
 from bitmind.utils import ExitContext, get_metadata
-from bitmind.wandb_utils import init_wandb
+from bitmind.wandb_utils import init_wandb, clean_wandb_cache
 from bitmind.types import CacheConfig, MediaType, Modality
 from bitmind.cache.sampler import ImageSampler
 from bitmind.generation import (
@@ -86,6 +86,8 @@ def __init__(self):
                     .metagraph(self.config.netuid)
                     .hotkeys.index(self.wallet.hotkey.ss58_address)
                 )
+                self.wandb_dir = str(Path(__file__).parent.parent)
+                clean_wandb_cache(self.wandb_dir)
                 self.wandb_run = init_wandb(
                     self.config.copy(),
                     self.config.wandb.process_name,
@@ -222,6 +224,7 @@ async def run(self):
                         if batch_count >= self.config.wandb.num_batches_per_run:
                             batch_count = 0
                             self.wandb_run.finish()
+                            clean_wandb_cache(self.wandb_dir)
                             self.wandb_run = init_wandb(
                                 self.config.copy(),
                                 self.config.wandb.process_name,