From 132f9853da197b48314692d5a2bfc646f9c72939 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Mon, 12 Oct 2026 12:01:52 -0500 Subject: [PATCH] fix(train): write a sourceable sm_training.env for quotes and non-identifier names The training container's bash script sources /opt/ml/input/sm_training.env under 'set -e' with an ERR trap. set_env wrote every variable as export NAME='value', so: - a single quote in any value (e.g. a hyperparameter 'it's', which also appears in SM_HPS and SM_TRAINING_ENV) broke the file's quoting, and - a channel name or hyperparameter key with '-' or '.' (allowed by the API, e.g. channel 'train-data' or key 'model.lr') produced an invalid identifier such as SM_CHANNEL_TRAIN-DATA, making 'export' fail, either of which killed the job before the entry script ran. Map non-identifier characters in names to '_' (as SM_HP_* already does for '-') and escape single quotes in values. Output for ordinary values is unchanged. Applied to the sagemaker-train and sagemaker-core copies. --- .../container_drivers/scripts/environment.py | 8 ++- .../unit/modules/train/test_environment.py | 59 +++++++++++++++++++ .../container_drivers/scripts/environment.py | 8 ++- .../scripts/test_enviornment.py | 49 +++++++++++++++ 4 files changed, 122 insertions(+), 2 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py b/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py index f48b8ac1a7..a8e5f4144f 100644 --- a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/scripts/environment.py +++ b/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) diff --git a/sagemaker-core/tests/unit/modules/train/test_environment.py b/sagemaker-core/tests/unit/modules/train/test_environment.py index 9f79133c63..2980ca950b 100644 --- a/sagemaker-core/tests/unit/modules/train/test_environment.py +++ b/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"] diff --git a/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py b/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py index f48b8ac1a7..a8e5f4144f 100644 --- a/sagemaker-train/src/sagemaker/train/container_drivers/scripts/environment.py +++ b/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) diff --git a/sagemaker-train/tests/unit/train/container_drivers/scripts/test_enviornment.py b/sagemaker-train/tests/unit/train/container_drivers/scripts/test_enviornment.py index 6a3cb0cf14..7d3fa50dd5 100644 --- a/sagemaker-train/tests/unit/train/container_drivers/scripts/test_enviornment.py +++ b/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"]