Utility functions for stage epochs (#1424) (#1433)

* Utility functions to get bound epochs per stage

* Handle start_epoch -1 case

* Post-process dict to remove -1's as possible

* Additional comments

* Outsource logic to helper function

* Iteratable bug fix

* nit - empty space

Author: KSGulin

src/sparseml/optim/manager.py

@@ -24,7 +24,7 @@
 from collections import OrderedDict
 from copy import deepcopy
 from functools import cmp_to_key
-from typing import Any, Dict, Generator, List, Optional, Union
+from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, Union
 
 from sparseml.optim.modifier import BaseModifier, BaseObject, ModifierProp
 from sparseml.sparsification.types import SparsificationTypes
@@ -343,47 +343,15 @@ def min_epochs(self) -> int:
         """
         :return: the minimum epochs required by any of the modifiers under the manager
         """
-        vals = []
-        vals.extend(
-            [
-                math.floor(mod.start_epoch)
-                for mod in self.iter_modifiers()
-                if mod.start_epoch > -1
-            ]
-        )
-        vals.extend(
-            [
-                math.floor(mod.end_epoch)
-                for mod in self.iter_modifiers()
-                if mod.end_epoch > -1
-            ]
-        )
-
-        return min(vals) if len(vals) > 0 else -1
+        return _min_modifier_epoch(self.iter_modifiers())
 
     @ModifierProp(serializable=False)
     def max_epochs(self) -> int:
         """
         :return: the maximum number of epochs required by any of the modifiers
             under the manager
         """
-        vals = []
-        vals.extend(
-            [
-                math.ceil(mod.start_epoch)
-                for mod in self.iter_modifiers()
-                if mod.start_epoch > -1
-            ]
-        )
-        vals.extend(
-            [
-                math.ceil(mod.end_epoch)
-                for mod in self.iter_modifiers()
-                if mod.end_epoch > -1
-            ]
-        )
-
-        return max(vals) if len(vals) > 0 else -1
+        return _max_modifier_epoch(self.iter_modifiers())
 
     def save(self, file_path: str, include_metadata: bool = True):
         """
@@ -561,6 +529,44 @@ def qat_active(self, epoch: float) -> bool:
             else False
         )
 
+    def get_start_end_epochs(self) -> Dict[str, Tuple[float, float]]:
+        """
+        Return an OrderedDict mapping each stage to its min and max epoch. If not a
+        staged manager, map 'all' to the the min and max epochs
+        """
+        if isinstance(self.modifiers, List):
+            return OrderedDict({"all": (self.min_epochs, self.max_epochs)})
+        else:
+            stage_max_min = OrderedDict()
+            for stage, mod_list in self.modifiers.items():
+                epoch_min = _min_modifier_epoch(mod_list)
+                epoch_max = _max_modifier_epoch(mod_list)
+                stage_max_min[stage] = (epoch_min, epoch_max)
+
+            # post-process to replace -1's with their real values
+            epochs_list = list(stage_max_min.values())
+            for i, (stage, epochs) in enumerate(stage_max_min.items()):
+                # replace start epochs that are -1 with the last epoch of the previous
+                # stage, or 0 if it's the first stage
+                if epochs[0] == -1:
+                    stage_max_min[stage][0] = epochs_list[i - 1][1] if i > 0 else 0
+                # replace end epochs that are -1 with the next stage's start epoch,
+                # unless it's the last stage
+                if epochs[1] == -1 and i < len(epochs_list) - 1:
+                    stage_max_min[stage][1] = epochs_list[i + 1][0]
+
+            return stage_max_min
+
+    def get_last_start_epoch(self) -> float:
+        """
+        Return the start epoch of the last stage in the recipe. Useful for applying
+        recipes at the correct epoch in a staged run
+        """
+        stage_max_min = self.get_start_end_epochs()
+        last_stage_epochs = stage_max_min[next(reversed(stage_max_min))]
+        last_start_epoch = last_stage_epochs[0]
+        return last_start_epoch if last_start_epoch > -1 else 0
+
     def _info_log_metadata(self):
         metadata_str = json.dumps(self._metadata, indent=1)
         _LOGGER.debug(f"Created recipe manager with metadata: {metadata_str}")
@@ -586,3 +592,28 @@ def _nested_dict_to_lines(
             # reached maximum nesting level.
             yaml_str_lines.append(indentation * nesting_depth + f"{key}: {value}")
     return yaml_str_lines
+
+
+def _min_modifier_epoch(modifiers: Iterable[BaseModifier]) -> float:
+    """
+    :return: the minimum epochs required by any of the modifiers provided
+    """
+    vals = [math.floor(mod.start_epoch) for mod in modifiers if mod.start_epoch > -1]
+
+    return min(vals) if len(vals) > 0 else -1
+
+
+def _max_modifier_epoch(modifiers: Iterable[BaseModifier]) -> float:
+    """
+    :return: the maximum number of epochs required by any of the modifiers provided
+    """
+    # save modifiers as list so it can iterated over multiple times
+    modifiers = [mod for mod in modifiers]
+
+    vals = []
+    vals.extend(
+        [math.ceil(mod.start_epoch) for mod in modifiers if mod.start_epoch > -1]
+    )
+    vals.extend([math.ceil(mod.end_epoch) for mod in modifiers if mod.end_epoch > -1])
+
+    return max(vals) if len(vals) > 0 else -1