R2 dataset integration

distributed_training/averaging/avg_handler.py

@@ -10,7 +10,7 @@
 from typing import Any, Dict, List, Tuple
 from distributed_training.averaging.exceptions import AllReduceError, ModelStateError
 from distributed_training.protocol import AllReduce
-from distributed_training.data.dataset import DatasetLoader
+from distributed_training.data.dataset_loader import DatasetLoader
 from distributed_training.utils.dendrite import (
     async_dendrite_forward,
 )
@@ -85,22 +85,22 @@ async def fetch_training_data(self, block):
         attempt = 0
         while attempt < self.retry_limit:
             try:
-                pages = await DatasetLoader.next_pages(
-                    offset=block,
-                    n_pages=5,
-                    seed=self.uid,
-                )
-                random.seed(self.uid)
-                random.shuffle(pages)
-
-                dataset = await DatasetLoader.create(
-                    batch_size=4,
-                    sequence_length=1024,
-                    pages_info=pages,
+                loader = DatasetLoader(
                     tokenizer=self.tokenizer,
+                    seed_base=self.uid,
+                    current_block=block,
                 )
 
-                return dataset
+                await loader.load_bucket_data_to_buffer()
+
+                # 1) add method="truncate" for debugging
+                # 2) default buffer quantity is 2300000 so we set to small quantity of 460000 (20%) 
+                #    cause here we only verify
+                loader.reduce_buffer_size(target_size=460000)                 
+
+                loader.prepare_batches()       
+
+                return loader                
             except Exception as e:
                 self.logger.error(f"Error fetching training data: {str(e)}")
                 attempt += 1

distributed_training/base/neuron.py

@@ -36,7 +36,7 @@
     R2Access,
     R2Config,
 )
-from distributed_training.utils.logger import setup_logging
+# from distributed_training.utils.logger import setup_logging
 from distributed_training.utils.misc import ttl_get_block
 from dotenv import load_dotenv
 
@@ -171,7 +171,7 @@ def __init__(self, config=None):
         self.master_uid = master_uid[0].item()
 
         # Setup Logging
-        setup_logging(self, config=self.config)
+        # setup_logging(self, config=self.config)
 
         # Create the R2 data model
         r2 = R2Config(