Override default save configs in xgboost (aws#341)

* Override default save configs in xgboost

* Change kwargs in SaveConfig.from_dict() to default_values

* default_values is an empty dict by default

* Remove unncessary dict copy

* Explicit is better than implicit

* Parse default_values in SaveConfigMode

Author: Edward J Kim

tornasole/core/json_config.py

@@ -80,7 +80,9 @@ def get_json_config_as_dict(json_config_path) -> Dict:
     return params_dict
 
 
-def create_hook_from_json_config(hook_cls, collection_manager, json_config_path):
+def create_hook_from_json_config(
+    hook_cls, collection_manager, json_config_path, default_values=None
+):
     """Returns a TornasoleHook object corresponding to either TF, PT, or MXNet.
 
     If json_config_path is None, an environment variable must be set.
@@ -100,7 +102,7 @@ def create_hook_from_json_config(hook_cls, collection_manager, json_config_path)
     out_dir = tornasole_params.get("out_dir", DEFAULT_SAGEMAKER_TORNASOLE_PATH)
     dry_run = tornasole_params.get("dry_run", False)
     reduction_config = tornasole_params.get(TORNASOLE_CONFIG_RDN_CFG_KEY)
-    save_config = SaveConfig.from_dict(tornasole_params.get("save_config_modes"))
+    save_config = SaveConfig.from_dict(tornasole_params.get("save_config_modes"), default_values)
     include_regex = tornasole_params.get(TORNASOLE_CONFIG_INCLUDE_REGEX_KEY)
     save_all = tornasole_params.get(TORNASOLE_CONFIG_SAVE_ALL_KEY, False)
     return hook_cls(

tornasole/core/save_config.py

@@ -108,7 +108,9 @@ def merge_default_save_config(self, default_save_config):
                     self.set_save_config(mode=mode, save_config_mode=SaveConfigMode())
 
     @classmethod
-    def from_dict(cls, params: Dict[ModeKeys, Any]) -> "SaveConfig":
+    def from_dict(
+        cls, params: Dict[ModeKeys, Any], default_values: Dict[str, Any] = None
+    ) -> "SaveConfig":
         """Parses a dict into a SaveConfig object.
 
     Appropriate formats:
@@ -119,12 +121,17 @@ def from_dict(cls, params: Dict[ModeKeys, Any]) -> "SaveConfig":
     """
         if params is None:
             return None
+        if default_values is None:
+            default_values = {}
         # Maybe convert strings to enums
-        if all([isinstance(key, str) for key, value in params.items()]):
+        if all(isinstance(key, str) for key, value in params.items()):
             params = {ModeKeys[key]: value for key, value in params.items()}
         # Maybe convert dicts to SaveConfigMode
-        if all([value is None or isinstance(value, dict) for key, value in params.items()]):
-            params = {key: SaveConfigMode.from_dict(value) for key, value in params.items()}
+        if all(value is None or isinstance(value, dict) for key, value in params.items()):
+            params = {
+                key: SaveConfigMode.from_dict(value, default_values)
+                for key, value in params.items()
+            }
         return cls(mode_save_configs=params)
 
     @classmethod
@@ -171,18 +178,18 @@ def __repr__(self):
 
 class SaveConfigMode:
     """
-  Wrapping all the save configuration parameters into this object.
-  This would make it easier to set different save configuration for
-  different collections and for the base tensors saved.
+    Wrapping all the save configuration parameters into this object.
+    This would make it easier to set different save configuration for
+    different collections and for the base tensors saved.
 
-  This class should not be serialized by itself, only inside of SaveConfig.
+    This class should not be serialized by itself, only inside of SaveConfig.
 
-  Parameters:
-    save_interval (int): Save every n steps.
-    save_steps (list of int): Save at all the steps given in this list. Overrides save_interval.
-    start_step (int): Save after n steps.
-    end_step (int): Stop saving after n steps.
-  """
+    Parameters:
+      save_interval (int): Save every n steps.
+      save_steps (list of int): Save at all the steps given in this list. Overrides save_interval.
+      start_step (int): Save after n steps.
+      end_step (int): Stop saving after n steps.
+    """
 
     def __init__(
         self,
@@ -222,16 +229,20 @@ def to_json_dict(self):
         }
 
     @classmethod
-    def from_dict(cls, params: Dict[str, Any]):
+    def from_dict(cls, params: Dict[str, Any], default_values: Dict[str, Any] = None):
         if params is None:
             return None
         elif not isinstance(params, dict):
             raise TypeError(f"params={params} is not a dict.")
+        if default_values is None:
+            default_values = {}
+        elif not isinstance(default_values, dict):
+            raise TypeError(f"default_values={default_values} is not a dict.")
         return cls(
-            save_interval=params.get("save_interval"),
-            start_step=params.get("start_step"),
-            end_step=params.get("end_step"),
-            save_steps=params.get("save_steps"),
+            save_interval=params.get("save_interval", default_values.get("save_interval")),
+            start_step=params.get("start_step", default_values.get("start_step")),
+            end_step=params.get("end_step", default_values.get("end_step")),
+            save_steps=params.get("save_steps", default_values.get("save_steps")),
         )
 
     def __eq__(self, other):

tornasole/xgboost/hook.py

@@ -9,10 +9,7 @@
 from tornasole.core.hook import CallbackHook
 from tornasole.core.tfevent.util import make_numpy_array
 from tornasole.core.access_layer.utils import training_has_ended
-from tornasole.core.json_config import (
-    create_hook_from_json_config,
-    TORNASOLE_CONFIG_DEFAULT_WORKER_NAME,
-)
+from tornasole.core.json_config import create_hook_from_json_config
 from tornasole.xgboost.singleton_utils import set_hook
 
 
@@ -21,8 +18,10 @@
 
 
 DEFAULT_INCLUDE_COLLECTIONS = [CollectionKeys.METRICS]
-
 DEFAULT_SAVE_CONFIG_INTERVAL = 10
+DEFAULT_SAVE_CONFIG_START_STEP = 0
+DEFAULT_SAVE_CONFIG_END_STEP = None
+DEFAULT_SAVE_CONFIG_SAVE_STEPS = []
 
 
 class TornasoleHook(CallbackHook):
@@ -120,8 +119,17 @@ def hook_from_config(cls, json_config_path=None):
         Otherwise,
             return None.
         """
+        default_values = dict(
+            save_interval=DEFAULT_SAVE_CONFIG_INTERVAL,
+            start_step=DEFAULT_SAVE_CONFIG_START_STEP,
+            end_step=DEFAULT_SAVE_CONFIG_END_STEP,
+            save_steps=DEFAULT_SAVE_CONFIG_SAVE_STEPS,
+        )
         return create_hook_from_json_config(
-            cls, get_collection_manager(), json_config_path=json_config_path
+            cls,
+            get_collection_manager(),
+            json_config_path=json_config_path,
+            default_values=default_values,
         )
 
     def _is_last_step(self, env: CallbackEnv) -> bool: