make data generation buffer_size configurable.

Author: hariharan-devarajan

dlio_benchmark/data_generator/indexed_binary_generator.py

@@ -55,9 +55,9 @@ def generate(self):
             sample_size = dim1 * dim2
             total_size = sample_size * self.num_samples
             write_size = total_size
-            MEMORY_SIZE = 2*GB
-            if total_size > MEMORY_SIZE:
-                write_size = MEMORY_SIZE - (MEMORY_SIZE % sample_size)
+            memory_size = self._args.generation_buffer_size
+            if total_size > memory_size:
+                write_size = memory_size - (memory_size % sample_size)
             out_path_spec = self.storage.get_uri(self._file_list[i])
             out_path_spec_off_idx = self.index_file_path_off(out_path_spec)
             out_path_spec_sz_idx = self.index_file_path_size(out_path_spec)

dlio_benchmark/utils/config.py

@@ -96,6 +96,7 @@ class ConfigArguments:
     do_eval: bool = False
     batch_size_eval: int = 1
     num_files_eval: int = 0
+    generation_buffer_size: int = 2 * 1073741824 # 2 GB
     eval_time: float = 0.0
     eval_time_stdev: float = 0.0
     eval_after_epoch: int = 1
@@ -348,6 +349,8 @@ def LoadConfig(args, config):
             args.num_files_train = config['dataset']['num_files_train']
         if 'num_files_eval' in config['dataset']:
             args.num_files_eval = config['dataset']['num_files_eval']
+        if 'generation_buffer_size' in config['dataset']:
+            args.generation_buffer_size = config['dataset']['generation_buffer_size']
         if 'num_samples_per_file' in config['dataset']:
             args.num_samples_per_file = config['dataset']['num_samples_per_file']
         if 'data_folder' in config['dataset']: