refactor enum for better naming

Author: hariharan-devarajan

dlio_benchmark/common/enumerations.py

@@ -17,12 +17,12 @@
 
 from enum import Enum
 
-class CheckpointType(Enum):
+class CheckpointLocationType(Enum):
     """
     Different types of underlying storage
     """
-    COLLECTIVE = 'collective'
-    INDEPENDENT = 'independent'
+    RANK_ZERO = 'rank_zero'
+    ALL_RANKS = 'all_ranks'
 
     def __str__(self):
         return self.value

dlio_benchmark/configs/workload/megatron_deepspeed.yaml

@@ -29,7 +29,7 @@ checkpoint:
   checkpoint_folder: checkpoints/megatron-deepspeed
   steps_between_checkpoints: 1000
   model_size: 30102
-  type: independent
+  type: all_ranks
   optimization_groups: [1009254400, 865075200, 793600]
   num_layers: 44
   layer_parameters: [129761280, 20971520]

dlio_benchmark/framework/tf_framework.py

@@ -29,7 +29,7 @@
 from dlio_benchmark.profiler.profiler_factory import ProfilerFactory
 from dlio_benchmark.storage.storage_factory import StorageFactory
 from dlio_benchmark.common.enumerations import FrameworkType, Profiler, FormatType, DatasetType, MetadataType, \
-    DataLoaderType, CheckpointType
+    DataLoaderType, CheckpointLocationType
 
 import tensorflow as tf
 from tensorflow.python.framework import errors
@@ -55,11 +55,11 @@ def __init__(self, profiling):
         self.reader_handler = None
         self.model_state = None
         rank_to_checkpoint = self.args.my_rank
-        if self.args.checkpoint_type == CheckpointType.COLLECTIVE:
+        if self.args.checkpoint_type == CheckpointLocationType.RANK_ZERO:
             rank_to_checkpoint = 0
         if rank_to_checkpoint == self.args.my_rank:
             num_ranks = 1
-            if self.args.checkpoint_type == CheckpointType.COLLECTIVE:
+            if self.args.checkpoint_type == CheckpointLocationType.RANK_ZERO:
                 num_ranks = self.args.comm_size
             if self.args.model_size > 0:
                 self.model_state = {"a": self._get_tensor(self.args.model_size*num_ranks)}
@@ -121,7 +121,7 @@ def checkpoint(self, epoch, step_number):
         """
         my_rank = DLIOMPI.get_instance().rank()
         rank_to_checkpoint = my_rank
-        if self.args.checkpoint_type == CheckpointType.COLLECTIVE:
+        if self.args.checkpoint_type == CheckpointLocationType.RANK_ZERO:
             rank_to_checkpoint = 0
         if rank_to_checkpoint == my_rank:
             if self.model_state:

dlio_benchmark/framework/torch_framework.py

@@ -16,7 +16,7 @@
 """
 
 from dlio_benchmark.common.error_code import ErrorCodes
-from dlio_benchmark.common.enumerations import FormatType, FrameworkType, DatasetType, DataLoaderType, CheckpointType
+from dlio_benchmark.common.enumerations import FormatType, FrameworkType, DatasetType, DataLoaderType, CheckpointLocationType
 from dlio_benchmark.data_loader.data_loader_factory import DataLoaderFactory
 from dlio_benchmark.framework.framework import Framework, DummyTraceObject
 from dlio_benchmark.common.constants import MODULE_AI_FRAMEWORK
@@ -62,7 +62,7 @@ def __init__(self, profiling):
         self.profiling = profiling
         self.reader_handler = None
         rank_to_checkpoint = self.args.my_rank
-        if self.args.checkpoint_type == CheckpointType.COLLECTIVE:
+        if self.args.checkpoint_type == CheckpointLocationType.RANK_ZERO:
             rank_to_checkpoint = 0
         if rank_to_checkpoint == self.args.my_rank:
             self.model_state = None
@@ -120,7 +120,7 @@ def trace_object(self, string, step, r):
     def checkpoint(self, epoch, step_number):
 
         rank_to_checkpoint = DLIOMPI.get_instance().rank()
-        if self.args.checkpoint_type == CheckpointType.COLLECTIVE:
+        if self.args.checkpoint_type == CheckpointLocationType.RANK_ZERO:
             rank_to_checkpoint = 0
         if rank_to_checkpoint == DLIOMPI.get_instance().rank():
             my_rank = DLIOMPI.get_instance().rank()

dlio_benchmark/utils/config.py

@@ -26,7 +26,8 @@
 
 from dlio_benchmark.common.constants import MODULE_CONFIG
 from dlio_benchmark.common.enumerations import StorageType, FormatType, Shuffle, ReadType, FileAccess, Compression, \
-    FrameworkType, DataLoaderType, Profiler, DatasetType, DataLoaderSampler, CheckpointType
+    FrameworkType, \
+    DataLoaderType, Profiler, DatasetType, DataLoaderSampler, CheckpointLocationType
 from dlio_benchmark.utils.utility import DLIOMPI
 from dataclasses import dataclass
 import math
@@ -99,7 +100,7 @@ class ConfigArguments:
     eval_time_stdev: float = 0.0
     eval_after_epoch: int = 1
     epochs_between_evals: int = 1
-    checkpoint_type: CheckpointType = CheckpointType.COLLECTIVE
+    checkpoint_type: CheckpointLocationType = CheckpointLocationType.RANK_ZERO
     model_size: int = 10240
     optimization_groups:  ClassVar[List[int]] = []
     num_layers: int = 1
@@ -453,7 +454,7 @@ def LoadConfig(args, config):
         if 'steps_between_checkpoints' in config['checkpoint']:
             args.steps_between_checkpoints = config['checkpoint']['steps_between_checkpoints']
         if 'type' in config['checkpoint']:
-            args.checkpoint_type = CheckpointType(config['checkpoint']['type'])
+            args.checkpoint_type = CheckpointLocationType(config['checkpoint']['type'])
         if 'model_size' in config['checkpoint']:
             args.model_size = config['checkpoint']['model_size']
         if 'optimization_groups' in config['checkpoint']:

tests/dlio_benchmark_test.py

@@ -209,12 +209,12 @@ def test_iostat_profiling() -> None:
         clean()
 
 @pytest.mark.timeout(60, method="thread")
-@pytest.mark.parametrize("framework, model_size, optimizers, num_layers, layer_params, type", [("tensorflow", 1024, [1024, 128], 2, [16], "independent"),
-                                                                                         ("pytorch", 1024, [1024, 128], 2, [16], "independent"),
-                                                                                         ("tensorflow", 1024, [1024, 128], 2, [16], "collective"),
-                                                                                         ("pytorch", 1024, [1024, 128], 2, [16], "collective"),
-                                                                                         ("tensorflow", 1024, [128], 1, [], "independent"),
-                                                                                         ("pytorch", 1024, [128], 1, [], "independent")])
+@pytest.mark.parametrize("framework, model_size, optimizers, num_layers, layer_params, type", [("tensorflow", 1024, [1024, 128], 2, [16], "all_ranks"),
+                                                                                         ("pytorch", 1024, [1024, 128], 2, [16], "all_ranks"),
+                                                                                         ("tensorflow", 1024, [1024, 128], 2, [16], "rank_zero"),
+                                                                                         ("pytorch", 1024, [1024, 128], 2, [16], "rank_zero"),
+                                                                                         ("tensorflow", 1024, [128], 1, [], "all_ranks"),
+                                                                                         ("pytorch", 1024, [128], 1, [], "all_ranks")])
 def test_checkpoint_epoch(framework, model_size, optimizers, num_layers, layer_params, type) -> None:
     clean()
     if comm.rank == 0:
@@ -249,7 +249,7 @@ def test_checkpoint_epoch(framework, model_size, optimizers, num_layers, layer_p
         if len(layer_params) > 0:
             n = num_layers
         nranks = 1
-        if type == "independent":
+        if type == "all_ranks":
             nranks = comm.size
         if framework == "tensorflow":
             num_check_files = 8 / 2 * (2 + 2 + 2*n) * nranks + 1