Mask Sensitive Env Logs in Container (#1568)

Author: benieric

src/sagemaker/modules/templates.py

@@ -73,10 +73,6 @@
 cat /opt/ml/input/config/inputdataconfig.json
 echo
 
-echo "/opt/ml/input/config/hyperparameters.json:"
-cat /opt/ml/input/config/hyperparameters.json
-echo
-
 echo "/opt/ml/input/data/sm_drivers/sourcecodeconfig.json"
 cat /opt/ml/input/data/sm_drivers/sourcecodeconfig.json
 echo

src/sagemaker/modules/testing_notebooks/base_model_trainer.ipynb

@@ -74,16 +74,24 @@
     "\n",
     "pytorch_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
     "\n",
-    "pytorch_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
-    "\n",
     "source_code_config = SourceCodeConfig(\n",
     "    source_dir=\"basic-script-mode\",\n",
     "    command=\"python custom_script.py\",\n",
     ")\n",
     "\n",
+    "hyperparameters = {\n",
+    "    \"secret_token\": \"123456\",\n",
+    "}\n",
+    "\n",
+    "env_vars = {\n",
+    "    \"PASSWORD\": \"123456\"\n",
+    "}\n",
+    "\n",
     "model_trainer = ModelTrainer(\n",
     "    training_image=pytorch_image,\n",
     "    source_code_config=source_code_config,\n",
+    "    hyperparameters=hyperparameters,\n",
+    "    environment=env_vars,\n",
     ")\n",
     "\n",
     "model_trainer.train()"

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

@@ -47,6 +47,9 @@
 
 ENV_OUTPUT_FILE = "/opt/ml/input/data/sm_drivers/scripts/sm_training.env"
 
+SENSITIVE_KEYWORDS = ["SECRET", "PASSWORD", "KEY", "TOKEN", "PRIVATE", "CREDS", "CREDENTIALS"]
+HIDDEN_VALUE = "******"
+
 
 def num_cpus() -> int:
     """Return the number of CPUs available in the current container.
@@ -199,6 +202,50 @@ def set_env(
             else:
                 f.write(f"export {key}='{value}'\n")
 
+    logger.info("Environment Variables:")
+    log_env_variables(env_vars_dict=env_vars)
+
+
+def mask_sensitive_info(data):
+    """Recursively mask sensitive information in a dictionary."""
+    if isinstance(data, dict):
+        for k, v in data.items():
+            if isinstance(v, dict):
+                data[k] = mask_sensitive_info(v)
+            elif isinstance(v, str) and any(
+                keyword.lower() in k.lower() for keyword in SENSITIVE_KEYWORDS
+            ):
+                data[k] = HIDDEN_VALUE
+    return data
+
+
+def log_key_value(key: str, value: str):
+    """Log a key-value pair, masking sensitive values if necessary."""
+    if any(keyword.lower() in key.lower() for keyword in SENSITIVE_KEYWORDS):
+        logger.info("%s=%s", key, HIDDEN_VALUE)
+    elif isinstance(value, dict):
+        masked_value = mask_sensitive_info(value)
+        logger.info("%s=%s", key, json.dumps(masked_value))
+    else:
+        try:
+            decoded_value = json.loads(value)
+            if isinstance(decoded_value, dict):
+                masked_value = mask_sensitive_info(decoded_value)
+                logger.info("%s=%s", key, json.dumps(masked_value))
+            else:
+                logger.info("%s=%s", key, decoded_value)
+        except (json.JSONDecodeError, TypeError):
+            logger.info("%s=%s", key, value)
+
+
+def log_env_variables(env_vars_dict: Dict[str, Any]):
+    """Log Environment Variables from the environment and an env_vars_dict."""
+    for key, value in os.environ.items():
+        log_key_value(key, value)
+
+    for key, value in env_vars_dict.items():
+        log_key_value(key, value)
+
 
 def main():
     """Main function to set the environment variables for the training job container."""

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

@@ -14,10 +14,18 @@
 from __future__ import absolute_import
 
 import os
+import io
+import logging
 
 from unittest.mock import patch
 
-from sagemaker.modules.train.container_drivers.scripts.environment import set_env
+from sagemaker.modules.train.container_drivers.scripts.environment import (
+    set_env,
+    log_key_value,
+    log_env_variables,
+    mask_sensitive_info,
+    HIDDEN_VALUE,
+)
 
 RESOURCE_CONFIG = dict(
     current_host="algo-1",
@@ -129,6 +137,39 @@ def test_set_env(mock_num_cpus, mock_num_gpus, mock_num_neurons):
         assert not os.path.exists(OUTPUT_FILE)
 
 
+@patch.dict(os.environ, {"SECRET_TOKEN": "122345678", "CLEAR_DATA": "123456789"}, clear=True)
+def test_log_env_variables():
+    log_stream = io.StringIO()
+    handler = logging.StreamHandler(log_stream)
+
+    logger = logging.getLogger("sagemaker.modules.train.container_drivers.scripts.environment")
+    logger.addHandler(handler)
+    logger.setLevel(logging.INFO)
+
+    env_vars = {
+        "SM_MODEL_DIR": "/opt/ml/model",
+        "SM_INPUT_DIR": "/opt/ml/input",
+        "SM_HPS": {"batch_size": 32, "learning_rate": 0.001, "access_token": "123456789"},
+        "SM_HP_BATCH_SIZE": 32,
+        "SM_HP_LEARNING_RATE": 0.001,
+        "SM_HP_ACCESS_TOKEN": "123456789",
+    }
+    log_env_variables(env_vars_dict=env_vars)
+
+    log_output = log_stream.getvalue()
+
+    assert f"SECRET_TOKEN={HIDDEN_VALUE}" in log_output
+    assert "CLEAR_DATA=123456789" in log_output
+    assert "SM_MODEL_DIR=/opt/ml/model" in log_output
+    assert (
+        f'SM_HPS={{"batch_size": 32, "learning_rate": 0.001, "access_token": "{HIDDEN_VALUE}"}}'
+        in log_output
+    )
+    assert "SM_HP_BATCH_SIZE=32" in log_output
+    assert "SM_HP_LEARNING_RATE=0.001" in log_output
+    assert f"SM_HP_ACCESS_TOKEN={HIDDEN_VALUE}" in log_output
+
+
 def _remove_extra_lines(string):
     """Removes extra blank lines from a string."""
     return "\n".join([line for line in string.splitlines() if line.strip()])