fix(train): write a sourceable sm_training.env for quotes and non-identifier names

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

@@ -19,6 +19,7 @@
 import subprocess
 import json
 import os
+import re
 import sys
 from pathlib import Path
 import logging
@@ -238,7 +239,12 @@ def set_env(
     }
     with open(output_file, "w") as f:
         for key, value in env_vars.items():
-            f.write(f"export {key}='{safe_serialize(value)}'\n")
+            # The file is sourced by a bash script running with 'set -e', so names must be valid
+            # shell identifiers (channel names and hyperparameter keys may contain '-' or '.')
+            # and single quotes inside values must be escaped.
+            name = re.sub(r"[^A-Za-z0-9_]", "_", key)
+            quoted = safe_serialize(value).replace("'", "'\"'\"'")
+            f.write(f"export {name}='{quoted}'\n")
 
     logger.info("Environment Variables:")
     log_env_variables(env_vars_dict=env_vars)

sagemaker-core/tests/unit/modules/train/test_environment.py

@@ -383,3 +383,62 @@ def test_set_env_multiple_channels(
 
         # Verify file was written
         mock_file.assert_called_once()
+
+
+@patch(
+    "sagemaker.core.modules.train.container_drivers.scripts.environment.read_source_code_json",
+    return_value=None,
+)
+@patch(
+    "sagemaker.core.modules.train.container_drivers.scripts.environment.read_distributed_json",
+    return_value=None,
+)
+@patch(
+    "sagemaker.core.modules.train.container_drivers.scripts.environment.num_cpus", return_value=1
+)
+@patch(
+    "sagemaker.core.modules.train.container_drivers.scripts.environment.num_gpus", return_value=0
+)
+@patch(
+    "sagemaker.core.modules.train.container_drivers.scripts.environment.num_neurons", return_value=0
+)
+@patch("sagemaker.core.modules.train.container_drivers.scripts.environment.log_env_variables")
+def test_set_env_output_can_be_sourced_with_special_characters(*_mocks):
+    """The env file is sourced under 'set -e'; quotes and non-identifier names must not break it."""
+    import subprocess
+    import tempfile
+
+    from sagemaker.core.modules.train.container_drivers.scripts.environment import (
+        set_env as _set_env,
+    )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        output_file = os.path.join(tmp_dir, "sm_training.env")
+        with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}):
+            _set_env(
+                resource_config={
+                    "current_host": "algo-1",
+                    "current_instance_type": "ml.m5.xlarge",
+                    "hosts": ["algo-1"],
+                    "network_interface_name": "eth0",
+                },
+                input_data_config={"train-data": {"TrainingInputMode": "File"}},
+                hyperparameters_config={"model.lr": "0.1", "prompt": "it's a test"},
+                output_file=output_file,
+            )
+
+        result = subprocess.run(
+            [
+                "bash",
+                "-c",
+                'set -e; trap "exit 3" ERR; source "$1"; '
+                'printf "%s\n" "$SM_CHANNEL_TRAIN_DATA" "$SM_HP_MODEL_LR" "$SM_HP_PROMPT"',
+                "bash",
+                output_file,
+            ],
+            capture_output=True,
+            text=True,
+        )
+
+    assert result.returncode == 0, result.stderr
+    assert result.stdout.splitlines() == ["/opt/ml/input/data/train-data", "0.1", "it's a test"]

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

@@ -19,6 +19,7 @@
 import subprocess
 import json
 import os
+import re
 import sys
 from pathlib import Path
 import logging
@@ -238,7 +239,12 @@ def set_env(
     }
     with open(output_file, "w") as f:
         for key, value in env_vars.items():
-            f.write(f"export {key}='{safe_serialize(value)}'\n")
+            # The file is sourced by a bash script running with 'set -e', so names must be valid
+            # shell identifiers (channel names and hyperparameter keys may contain '-' or '.')
+            # and single quotes inside values must be escaped.
+            name = re.sub(r"[^A-Za-z0-9_]", "_", key)
+            quoted = safe_serialize(value).replace("'", "'\"'\"'")
+            f.write(f"export {name}='{quoted}'\n")
 
     logger.info("Environment Variables:")
     log_env_variables(env_vars_dict=env_vars)

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

@@ -212,3 +212,52 @@ def test_log_env_variables():
 def _remove_extra_lines(string):
     """Removes extra blank lines from a string."""
     return "\n".join([line for line in string.splitlines() if line.strip()])
+
+
+@patch(
+    "sagemaker.train.container_drivers.scripts.environment.read_source_code_json", return_value=None
+)
+@patch(
+    "sagemaker.train.container_drivers.scripts.environment.read_distributed_json", return_value=None
+)
+@patch("sagemaker.train.container_drivers.scripts.environment.num_cpus", return_value=1)
+@patch("sagemaker.train.container_drivers.scripts.environment.num_gpus", return_value=0)
+@patch("sagemaker.train.container_drivers.scripts.environment.num_neurons", return_value=0)
+@patch("sagemaker.train.container_drivers.scripts.environment.log_env_variables")
+def test_set_env_output_can_be_sourced_with_special_characters(*_mocks):
+    """The env file is sourced under 'set -e'; quotes and non-identifier names must not break it."""
+    import subprocess
+    import tempfile
+
+    from sagemaker.train.container_drivers.scripts.environment import set_env as _set_env
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        output_file = os.path.join(tmp_dir, "sm_training.env")
+        with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}):
+            _set_env(
+                resource_config={
+                    "current_host": "algo-1",
+                    "current_instance_type": "ml.m5.xlarge",
+                    "hosts": ["algo-1"],
+                    "network_interface_name": "eth0",
+                },
+                input_data_config={"train-data": {"TrainingInputMode": "File"}},
+                hyperparameters_config={"model.lr": "0.1", "prompt": "it's a test"},
+                output_file=output_file,
+            )
+
+        result = subprocess.run(
+            [
+                "bash",
+                "-c",
+                'set -e; trap "exit 3" ERR; source "$1"; '
+                'printf "%s\n" "$SM_CHANNEL_TRAIN_DATA" "$SM_HP_MODEL_LR" "$SM_HP_PROMPT"',
+                "bash",
+                output_file,
+            ],
+            capture_output=True,
+            text=True,
+        )
+
+    assert result.returncode == 0, result.stderr
+    assert result.stdout.splitlines() == ["/opt/ml/input/data/train-data", "0.1", "it's a test"]