Fix: Correctly serialize SM_HPS env var (#1611)

Author: benieric

src/sagemaker/modules/train/container_drivers/scripts/environment.py

@@ -24,7 +24,7 @@
 parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 sys.path.insert(0, parent_dir)
 
-from utils import safe_serialize  # noqa: E402 # pylint: disable=C0413
+from utils import safe_serialize, safe_deserialize  # noqa: E402 # pylint: disable=C0413
 
 # Initialize logger
 SM_LOG_LEVEL = os.environ.get("SM_LOG_LEVEL", 20)
@@ -115,6 +115,21 @@ def num_neurons() -> int:
         return 0
 
 
+def deserialize_hyperparameters(hyperparameters: Dict[str, str]) -> Dict[str, Any]:
+    """Deserialize hyperparameters from string to their original types.
+
+    Args:
+        hyperparameters (Dict[str, str]): Hyperparameters as strings.
+
+    Returns:
+        Dict[str, Any]: Hyperparameters as their original types.
+    """
+    deserialized_hyperparameters = {}
+    for key, value in hyperparameters.items():
+        deserialized_hyperparameters[key] = safe_deserialize(value)
+    return deserialized_hyperparameters
+
+
 def set_env(
     resource_config: Dict[str, Any],
     input_data_config: Dict[str, Any],
@@ -150,10 +165,11 @@ def set_env(
     env_vars["SM_CHANNELS"] = channels
 
     # Hyperparameters
-    env_vars["SM_HPS"] = hyperparameters_config
-    for key, value in hyperparameters_config.items():
+    hps = deserialize_hyperparameters(hyperparameters_config)
+    for key, value in hps.items():
         key_upper = key.replace("-", "_").upper()
-        env_vars[f"SM_HP_{key_upper}"] = safe_serialize(value)
+        env_vars[f"SM_HP_{key_upper}"] = value
+    env_vars["SM_HPS"] = hps
 
     # Host Variables
     current_host = resource_config["current_host"]

src/sagemaker/modules/train/model_trainer.py

@@ -119,7 +119,7 @@ class ModelTrainer(BaseModel):
     ```
 
     Attributes:
-        session (Optiona(Session)):
+        sagemaker_session (Optiona(Session)):
             The SageMakerCore session. For convinience, can be imported like:
             `from sagemaker.modules import Session`.
             If not specified, a new session will be created.

tests/data/modules/params_script/train.py

@@ -99,6 +99,43 @@ def main():
     assert json.loads(os.environ["SM_HP_LIST"]) == EXPECTED_HYPERPARAMETERS["list"]
     assert json.loads(os.environ["SM_HP_DICT"]) == EXPECTED_HYPERPARAMETERS["dict"]
 
+    params = json.loads(os.environ["SM_HPS"])
+    print(f"SM_HPS: {params}")
+    assert params["string"] == EXPECTED_HYPERPARAMETERS["string"]
+    assert params["integer"] == EXPECTED_HYPERPARAMETERS["integer"]
+    assert params["boolean"] == EXPECTED_HYPERPARAMETERS["boolean"]
+    assert params["float"] == EXPECTED_HYPERPARAMETERS["float"]
+    assert params["list"] == EXPECTED_HYPERPARAMETERS["list"]
+    assert params["dict"] == EXPECTED_HYPERPARAMETERS["dict"]
+
+    assert isinstance(params, dict)
+    assert isinstance(params["string"], str)
+    assert isinstance(params["integer"], int)
+    assert isinstance(params["boolean"], bool)
+    assert isinstance(params["float"], float)
+    assert isinstance(params["list"], list)
+    assert isinstance(params["dict"], dict)
+
+    params = json.loads(os.environ["SM_TRAINING_ENV"])["hyperparameters"]
+    print(params)
+    assert params["string"] == EXPECTED_HYPERPARAMETERS["string"]
+    assert params["integer"] == EXPECTED_HYPERPARAMETERS["integer"]
+    assert params["boolean"] == EXPECTED_HYPERPARAMETERS["boolean"]
+    assert params["float"] == EXPECTED_HYPERPARAMETERS["float"]
+    assert params["list"] == EXPECTED_HYPERPARAMETERS["list"]
+    assert params["dict"] == EXPECTED_HYPERPARAMETERS["dict"]
+
+    assert isinstance(params, dict)
+    assert isinstance(params["string"], str)
+    assert isinstance(params["integer"], int)
+    assert isinstance(params["boolean"], bool)
+    assert isinstance(params["float"], float)
+    assert isinstance(params["list"], list)
+    assert isinstance(params["dict"], dict)
+    print(f"SM_TRAINING_ENV -> hyperparameters: {params}")
+
+    print("Test passed.")
+
 
 if __name__ == "__main__":
     main()

tests/unit/sagemaker/modules/train/container_drivers/scripts/test_enviornment.py

@@ -26,7 +26,7 @@
     mask_sensitive_info,
     HIDDEN_VALUE,
 )
-from sagemaker.modules.train.container_drivers.utils import safe_serialize
+from sagemaker.modules.train.container_drivers.utils import safe_serialize, safe_deserialize
 
 RESOURCE_CONFIG = dict(
     current_host="algo-1",
@@ -92,11 +92,11 @@
 export SM_CHANNEL_TRAIN='/opt/ml/input/data/train'
 export SM_CHANNEL_VALIDATION='/opt/ml/input/data/validation'
 export SM_CHANNELS='["train", "validation"]'
-export SM_HPS='{"batch_size": 32, "learning_rate": 0.001, "hosts": ["algo-1", "algo-2"], "mp_parameters": {"microbatches": 2, "partitions": 2, "pipeline": "interleaved", "optimize": "memory", "horovod": true}}'
 export SM_HP_BATCH_SIZE='32'
 export SM_HP_LEARNING_RATE='0.001'
 export SM_HP_HOSTS='["algo-1", "algo-2"]'
 export SM_HP_MP_PARAMETERS='{"microbatches": 2, "partitions": 2, "pipeline": "interleaved", "optimize": "memory", "horovod": true}'
+export SM_HPS='{"batch_size": 32, "learning_rate": 0.001, "hosts": ["algo-1", "algo-2"], "mp_parameters": {"microbatches": 2, "partitions": 2, "pipeline": "interleaved", "optimize": "memory", "horovod": true}}'
 export SM_CURRENT_HOST='algo-1'
 export SM_CURRENT_INSTANCE_TYPE='ml.p3.16xlarge'
 export SM_HOSTS='["algo-1", "algo-2", "algo-3"]'
@@ -119,7 +119,13 @@
     "sagemaker.modules.train.container_drivers.scripts.environment.safe_serialize",
     side_effect=safe_serialize,
 )
-def test_set_env(mock_safe_serialize, mock_num_cpus, mock_num_gpus, mock_num_neurons):
+@patch(
+    "sagemaker.modules.train.container_drivers.scripts.environment.safe_deserialize",
+    side_effect=safe_deserialize,
+)
+def test_set_env(
+    mock_safe_deserialize, mock_safe_serialize, mock_num_cpus, mock_num_gpus, mock_num_neurons
+):
     with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}):
         set_env(
             resource_config=RESOURCE_CONFIG,