From 7fde2ec459a9173e32143e9c6e6fb624d1b8290c Mon Sep 17 00:00:00 2001 From: Jaya Kasirajan Date: Mon, 12 Oct 2026 23:10:22 +0000 Subject: [PATCH] fix(train): shell-quote values written to the training env file sm_training.env is consumed with `source`, so a value containing a single quote ended the quoting and the remainder ran as shell. A legitimate value such as "it's 0.001" was also silently truncated. Quote each value with shlex.quote. --- .../container_drivers/scripts/environment.py | 4 +- .../bootstrap_runtime_environment.py | 4 +- .../unit/modules/train/test_environment.py | 63 ++++++++++ .../test_bootstrap_runtime_environment.py | 43 +++++++ .../container_drivers/scripts/environment.py | 4 +- .../bootstrap_runtime_environment.py | 4 +- .../scripts/test_enviornment.py | 110 ++++++++++++++---- .../test_bootstrap_runtime_environment.py | 41 +++++++ 8 files changed, 244 insertions(+), 29 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..cdc510d65c 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 shlex import sys from pathlib import Path import logging @@ -238,7 +239,8 @@ 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 consumed with `source`, so the value must be shell-quoted. + f.write(f"export {key}={shlex.quote(safe_serialize(value))}\n") logger.info("Environment Variables:") log_env_variables(env_vars_dict=env_vars) diff --git a/sagemaker-core/src/sagemaker/core/remote_function/runtime_environment/bootstrap_runtime_environment.py b/sagemaker-core/src/sagemaker/core/remote_function/runtime_environment/bootstrap_runtime_environment.py index 841995af10..569739274f 100644 --- a/sagemaker-core/src/sagemaker/core/remote_function/runtime_environment/bootstrap_runtime_environment.py +++ b/sagemaker-core/src/sagemaker/core/remote_function/runtime_environment/bootstrap_runtime_environment.py @@ -20,6 +20,7 @@ import multiprocessing import os import pathlib +import shlex import shutil import subprocess import sys @@ -527,7 +528,8 @@ 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 consumed with `source`, so the value must be shell-quoted. + f.write(f"export {key}={shlex.quote(safe_serialize(value))}\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..1ab6a47658 100644 --- a/sagemaker-core/tests/unit/modules/train/test_environment.py +++ b/sagemaker-core/tests/unit/modules/train/test_environment.py @@ -13,6 +13,11 @@ import json import os +import shutil +import subprocess + +import pytest + from unittest.mock import patch, mock_open from sagemaker.core.modules.train.container_drivers.scripts.environment import ( num_cpus, @@ -383,3 +388,61 @@ def test_set_env_multiple_channels( # Verify file was written mock_file.assert_called_once() + + +_CORE_ENV = "sagemaker.core.modules.train.container_drivers.scripts.environment" + +_CORE_RESOURCE_CONFIG = { + "current_host": "algo-1", + "current_instance_type": "ml.m5.large", + "hosts": ["algo-1"], + "network_interface_name": "eth0", +} + + +def _core_write_env(hyperparameters, output_file): + """Run set_env with the given hyperparameters and no source code or distributed config.""" + patches = [ + patch(f"{_CORE_ENV}.read_source_code_json", return_value={}), + patch(f"{_CORE_ENV}.read_distributed_json", return_value={}), + patch(f"{_CORE_ENV}.num_cpus", return_value=1), + patch(f"{_CORE_ENV}.num_gpus", return_value=0), + patch(f"{_CORE_ENV}.num_neurons", return_value=0), + ] + with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}): + for p in patches: + p.start() + try: + set_env( + resource_config=_CORE_RESOURCE_CONFIG, + input_data_config={}, + hyperparameters_config=hyperparameters, + output_file=output_file, + ) + finally: + for p in patches: + p.stop() + + +@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file") +@pytest.mark.parametrize( + "payload", + ["it's 0.001'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"], +) +def test_set_env_hyperparameter_value_survives_source_verbatim(tmp_path, payload): + """A hyperparameter value must round-trip through `source` and never be executed.""" + marker = tmp_path / "marker" + payload = payload.format(marker=marker) + output_file = str(tmp_path / "sm_training.env") + + _core_write_env({"learning_rate": payload}, output_file) + + result = subprocess.run( + ["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_HP_LEARNING_RATE"'], + capture_output=True, + text=True, + check=True, + ) + + assert result.stdout == payload + assert not marker.exists() diff --git a/sagemaker-core/tests/unit/remote_function/runtime_environment/test_bootstrap_runtime_environment.py b/sagemaker-core/tests/unit/remote_function/runtime_environment/test_bootstrap_runtime_environment.py index d7b8dc152b..8a158f17e9 100644 --- a/sagemaker-core/tests/unit/remote_function/runtime_environment/test_bootstrap_runtime_environment.py +++ b/sagemaker-core/tests/unit/remote_function/runtime_environment/test_bootstrap_runtime_environment.py @@ -11,6 +11,10 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. +import os +import shutil +import subprocess + import pytest from unittest.mock import Mock, patch, mock_open @@ -543,3 +547,42 @@ def test_main_failure(self, mock_write_failure, mock_parse): assert exc_info.value.code == DEFAULT_FAILURE_CODE mock_write_failure.assert_called_once() + + +_BOOTSTRAP = "sagemaker.core.remote_function.runtime_environment.bootstrap_runtime_environment" + + +@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file") +@pytest.mark.parametrize( + "payload", + ["eth0'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"], +) +def test_set_env_value_survives_source_verbatim(tmp_path, payload): + """A value written to the env file must round-trip through `source` and never be executed.""" + marker = tmp_path / "marker" + payload = payload.format(marker=marker) + output_file = str(tmp_path / "sm_training.env") + resource_config = { + "current_host": "algo-1", + "current_instance_type": "ml.m5.large", + "hosts": ["algo-1"], + "network_interface_name": payload, + } + + with ( + patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}), + patch(f"{_BOOTSTRAP}.num_cpus", return_value=1), + patch(f"{_BOOTSTRAP}.num_gpus", return_value=0), + patch(f"{_BOOTSTRAP}.num_neurons", return_value=0), + ): + set_env(resource_config, output_file=output_file) + + result = subprocess.run( + ["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_NETWORK_INTERFACE_NAME"'], + capture_output=True, + text=True, + check=True, + ) + + assert result.stdout == payload + assert not marker.exists() 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..cdc510d65c 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 shlex import sys from pathlib import Path import logging @@ -238,7 +239,8 @@ 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 consumed with `source`, so the value must be shell-quoted. + f.write(f"export {key}={shlex.quote(safe_serialize(value))}\n") logger.info("Environment Variables:") log_env_variables(env_vars_dict=env_vars) diff --git a/sagemaker-train/src/sagemaker/train/remote_function/runtime_environment/bootstrap_runtime_environment.py b/sagemaker-train/src/sagemaker/train/remote_function/runtime_environment/bootstrap_runtime_environment.py index 18cf13d1f5..b9378d909c 100644 --- a/sagemaker-train/src/sagemaker/train/remote_function/runtime_environment/bootstrap_runtime_environment.py +++ b/sagemaker-train/src/sagemaker/train/remote_function/runtime_environment/bootstrap_runtime_environment.py @@ -20,6 +20,7 @@ import multiprocessing import os import pathlib +import shlex import shutil import subprocess import sys @@ -524,7 +525,8 @@ 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 consumed with `source`, so the value must be shell-quoted. + f.write(f"export {key}={shlex.quote(safe_serialize(value))}\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..cabf271e62 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 @@ -17,6 +17,10 @@ import os import io import logging +import shutil +import subprocess + +import pytest from unittest.mock import patch @@ -87,37 +91,37 @@ # flake8: noqa EXPECTED_ENV = """ -export SM_MODEL_DIR='/opt/ml/model' -export SM_INPUT_DIR='/opt/ml/input' -export SM_INPUT_DATA_DIR='/opt/ml/input/data' -export SM_INPUT_CONFIG_DIR='/opt/ml/input/config' -export SM_OUTPUT_DIR='/opt/ml/output' -export SM_OUTPUT_FAILURE='/opt/ml/output/failure' -export SM_OUTPUT_DATA_DIR='/opt/ml/output/data' -export SM_LOG_LEVEL='20' -export SM_MASTER_ADDR='algo-1' -export SM_MASTER_PORT='7777' -export SM_SOURCE_DIR='/opt/ml/input/data/code' -export SM_ENTRY_SCRIPT='train.py' -export SM_DISTRIBUTED_DRIVER_DIR='/opt/ml/input/data/sm_drivers/distributed_drivers' +export SM_MODEL_DIR=/opt/ml/model +export SM_INPUT_DIR=/opt/ml/input +export SM_INPUT_DATA_DIR=/opt/ml/input/data +export SM_INPUT_CONFIG_DIR=/opt/ml/input/config +export SM_OUTPUT_DIR=/opt/ml/output +export SM_OUTPUT_FAILURE=/opt/ml/output/failure +export SM_OUTPUT_DATA_DIR=/opt/ml/output/data +export SM_LOG_LEVEL=20 +export SM_MASTER_ADDR=algo-1 +export SM_MASTER_PORT=7777 +export SM_SOURCE_DIR=/opt/ml/input/data/code +export SM_ENTRY_SCRIPT=train.py +export SM_DISTRIBUTED_DRIVER_DIR=/opt/ml/input/data/sm_drivers/distributed_drivers export SM_DISTRIBUTED_CONFIG='{"process_count_per_node": 2}' -export SM_CHANNEL_TRAIN='/opt/ml/input/data/train' -export SM_CHANNEL_VALIDATION='/opt/ml/input/data/validation' +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_HP_BATCH_SIZE='32' -export SM_HP_LEARNING_RATE='0.001' +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_CURRENT_HOST=algo-1 +export SM_CURRENT_INSTANCE_TYPE=ml.p3.16xlarge export SM_HOSTS='["algo-1", "algo-2", "algo-3"]' -export SM_NETWORK_INTERFACE_NAME='eth0' -export SM_HOST_COUNT='3' -export SM_CURRENT_HOST_RANK='0' -export SM_NUM_CPUS='8' -export SM_NUM_GPUS='0' -export SM_NUM_NEURONS='0' +export SM_NETWORK_INTERFACE_NAME=eth0 +export SM_HOST_COUNT=3 +export SM_CURRENT_HOST_RANK=0 +export SM_NUM_CPUS=8 +export SM_NUM_GPUS=0 +export SM_NUM_NEURONS=0 export SM_RESOURCE_CONFIG='{"current_host": "algo-1", "hosts": ["algo-1", "algo-2", "algo-3"], "current_group_name": "train1", "current_instance_type": "ml.p3.16xlarge", "instance_groups": [{"instance_group_name": "train1", "instance_type": "ml.p3.16xlarge", "hosts": ["algo-1", "algo-2"]}, {"instance_group_name": "train2", "instance_type": "ml.p3.8xlarge", "hosts": ["algo-3"]}], "network_interface_name": "eth0"}' export SM_INPUT_DATA_CONFIG='{"train": {"ContentType": "trainingContentType", "TrainingInputMode": "File", "S3DistributionType": "FullyReplicated", "RecordWrapperType": "None"}, "validation": {"TrainingInputMode": "File", "S3DistributionType": "FullyReplicated", "RecordWrapperType": "None"}}' export SM_TRAINING_ENV='{"channel_input_dirs": {"train": "/opt/ml/input/data/train", "validation": "/opt/ml/input/data/validation"}, "current_host": "algo-1", "current_instance_type": "ml.p3.16xlarge", "hosts": ["algo-1", "algo-2", "algo-3"], "master_addr": "algo-1", "master_port": 7777, "hyperparameters": {"batch_size": 32, "learning_rate": 0.001, "hosts": ["algo-1", "algo-2"], "mp_parameters": {"microbatches": 2, "partitions": 2, "pipeline": "interleaved", "optimize": "memory", "horovod": true}}, "input_data_config": {"train": {"ContentType": "trainingContentType", "TrainingInputMode": "File", "S3DistributionType": "FullyReplicated", "RecordWrapperType": "None"}, "validation": {"TrainingInputMode": "File", "S3DistributionType": "FullyReplicated", "RecordWrapperType": "None"}}, "input_config_dir": "/opt/ml/input/config", "input_data_dir": "/opt/ml/input/data", "input_dir": "/opt/ml/input", "job_name": "test-job", "log_level": 20, "model_dir": "/opt/ml/model", "network_interface_name": "eth0", "num_cpus": 8, "num_gpus": 0, "num_neurons": 0, "output_data_dir": "/opt/ml/output/data", "resource_config": {"current_host": "algo-1", "hosts": ["algo-1", "algo-2", "algo-3"], "current_group_name": "train1", "current_instance_type": "ml.p3.16xlarge", "instance_groups": [{"instance_group_name": "train1", "instance_type": "ml.p3.16xlarge", "hosts": ["algo-1", "algo-2"]}, {"instance_group_name": "train2", "instance_type": "ml.p3.8xlarge", "hosts": ["algo-3"]}], "network_interface_name": "eth0"}}' @@ -212,3 +216,59 @@ 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()]) + + +_ENV_PATCHES = [ + patch( + "sagemaker.train.container_drivers.scripts.environment.read_source_code_json", + return_value={}, + ), + patch( + "sagemaker.train.container_drivers.scripts.environment.read_distributed_json", + return_value={}, + ), + 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), +] + + +def _write_env(hyperparameters, output_file): + """Run set_env with the given hyperparameters and no source code or distributed config.""" + with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}): + for p in _ENV_PATCHES: + p.start() + try: + set_env( + resource_config=RESOURCE_CONFIG, + input_data_config=INPUT_DATA_CONFIG, + hyperparameters_config=hyperparameters, + output_file=output_file, + ) + finally: + for p in _ENV_PATCHES: + p.stop() + + +@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file") +@pytest.mark.parametrize( + "payload", + ["it's 0.001'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"], +) +def test_set_env_hyperparameter_value_survives_source_verbatim(tmp_path, payload): + """A hyperparameter value must round-trip through `source` and never be executed.""" + marker = tmp_path / "marker" + payload = payload.format(marker=marker) + output_file = str(tmp_path / "sm_training.env") + + _write_env({"learning_rate": payload}, output_file) + + result = subprocess.run( + ["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_HP_LEARNING_RATE"'], + capture_output=True, + text=True, + check=True, + ) + + assert result.stdout == payload + assert not marker.exists() diff --git a/sagemaker-train/tests/unit/train/remote_function/test_bootstrap_runtime_environment.py b/sagemaker-train/tests/unit/train/remote_function/test_bootstrap_runtime_environment.py index 644fbe7a8d..7c293b57f4 100644 --- a/sagemaker-train/tests/unit/train/remote_function/test_bootstrap_runtime_environment.py +++ b/sagemaker-train/tests/unit/train/remote_function/test_bootstrap_runtime_environment.py @@ -15,6 +15,8 @@ from __future__ import absolute_import import json +import os +import shutil import pytest import subprocess from unittest.mock import patch, MagicMock, mock_open @@ -810,3 +812,42 @@ def test_main_non_root_user(self, mock_getuser, mock_manager_class): main(args) mock_manager.change_dir_permission.assert_called_once() + + +_BOOTSTRAP = "sagemaker.train.remote_function.runtime_environment.bootstrap_runtime_environment" + + +@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file") +@pytest.mark.parametrize( + "payload", + ["eth0'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"], +) +def test_set_env_value_survives_source_verbatim(tmp_path, payload): + """A value written to the env file must round-trip through `source` and never be executed.""" + marker = tmp_path / "marker" + payload = payload.format(marker=marker) + output_file = str(tmp_path / "sm_training.env") + resource_config = { + "current_host": "algo-1", + "current_instance_type": "ml.m5.large", + "hosts": ["algo-1"], + "network_interface_name": payload, + } + + with ( + patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}), + patch(f"{_BOOTSTRAP}.num_cpus", return_value=1), + patch(f"{_BOOTSTRAP}.num_gpus", return_value=0), + patch(f"{_BOOTSTRAP}.num_neurons", return_value=0), + ): + set_env(resource_config, output_file=output_file) + + result = subprocess.run( + ["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_NETWORK_INTERFACE_NAME"'], + capture_output=True, + text=True, + check=True, + ) + + assert result.stdout == payload + assert not marker.exists()