add more checks for model base (#12)

Co-authored-by: Guoshuai Zhao <guzhao@microsoft.com>

superbench/benchmarks/model_benchmarks/model_base.py

@@ -39,10 +39,12 @@ def __init__(self, name, parameters=''):
         super().__init__(name, parameters)
 
         self._benchmark_type = BenchmarkType.MODEL
-        self._world_size = None
+        self._world_size = 1
+        self._local_rank = None
         self._dataset = None
         self._dataloader = None
         self._model = None
+        self._optimizer_type = None
         self._optimizer = None
         self._loss_fn = None
         self._target = None
@@ -107,17 +109,29 @@ def add_parser_arguments(self):
 
     @abstractmethod
     def _init_distributed_setting(self):
-        """Initialize the distributed library and bind the worker to GPU."""
+        """Initialize the distributed library and bind the worker to GPU.
+
+        Return:
+            True if distributed library is initialized successfully.
+        """
         pass
 
     @abstractmethod
     def _generate_dataset(self):
-        """Generate dataset for benchmarking according to shape info."""
+        """Generate dataset for benchmarking according to shape info.
+
+        Return:
+            True if dataset is created successfully.
+        """
         pass
 
     @abstractmethod
     def _init_dataloader(self):
-        """Initialize the distributed dataloader."""
+        """Initialize the dataloader.
+
+        Return:
+            True if dataloader is created successfully.
+        """
         pass
 
     def _preprocess(self):
@@ -126,18 +140,30 @@ def _preprocess(self):
         Return:
             True if _preprocess() succeed.
         """
-        ret = super()._preprocess()
-        if not ret:
+        if not super()._preprocess():
+            return False
+
+        if not self._init_distributed_setting():
+            self._result.set_return_code(ReturnCode.DISTRIBUTED_SETTING_INIT_FAILURE)
+            return False
+
+        if not self._generate_dataset():
+            self._result.set_return_code(ReturnCode.DATASET_GENERATION_FAILURE)
+            return False
+
+        if not self._init_dataloader():
+            self._result.set_return_code(ReturnCode.DATALOADER_INIT_FAILURE)
             return False
 
-        self._init_distributed_setting()
-        self._generate_dataset()
-        self._init_dataloader()
         return True
 
     @abstractmethod
     def _create_optimizer(self):
-        """Create the optimzier instance used for training."""
+        """Create the optimzier instance used for training and wrap with distributed library if need.
+
+        Return:
+            True if optimizer instance is created successfully.
+        """
         pass
 
     @abstractmethod
@@ -158,8 +184,14 @@ def __train(self, precision):
         Return:
             True if step_times list is not empty.
         """
-        self._create_model(precision)
-        self._create_optimizer()
+        if not self._create_model(precision):
+            self._result.set_return_code(ReturnCode.MODEL_CREATION_FAILURE)
+            return False
+
+        if not self._create_optimizer():
+            self._result.set_return_code(ReturnCode.OPTIMIZER_CREATION_FAILURE)
+            return False
+
         # The unit of step time should be millisecond.
         step_times = self._train_step(precision)
         if len(step_times) == 0:

superbench/benchmarks/return_code.py

 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-"""A module for unified context of benchmarks."""
+"""A module for statuses of benchmarks."""
 
-import enum
-
-
-class Enum(enum.Enum):
-    """Customized Enum class."""
-    @classmethod
-    def get_values(cls):
-        """Return the value list."""
-        values = [item.value for item in cls]
-        return values
+from superbench.benchmarks.context import Enum
 
 
 class ReturnCode(Enum):
@@ -26,3 +17,8 @@ class ReturnCode(Enum):
     NO_SUPPORTED_PRECISION = 10
     MODEL_TRAIN_FAILURE = 11
     MODEL_INFERENCE_FAILURE = 12
+    DISTRIBUTED_SETTING_INIT_FAILURE = 13
+    DATASET_GENERATION_FAILURE = 14
+    DATALOADER_INIT_FAILURE = 15
+    OPTIMIZER_CREATION_FAILURE = 16
+    MODEL_CREATION_FAILURE = 17

tests/benchmarks/model_benchmarks/test_model_base.py

@@ -41,23 +41,23 @@ def add_parser_arguments(self):
 
     def _init_distributed_setting(self):
         """Initialize the distributed library and bind the worker to GPU."""
-        pass
+        return True
 
     def _generate_dataset(self):
         """Generate dataset for benchmarking according to shape info."""
-        pass
+        return True
 
     def _init_dataloader(self):
         """Initialize the distributed dataloader."""
-        pass
+        return True
 
     def _create_optimizer(self):
         """Create the optimzier instance used for training."""
-        pass
+        return True
 
     def _create_model(self, precision):
         """Construct the model for benchmarking."""
-        pass
+        return True
 
     def _train_step(self, precision):
         """Define the training process.