Skip to content

Commit 8d07fea

Browse files
jkasirajJaya Kasirajan
jkasiraj
and
Jaya Kasirajan
authored
fix(train): shell-quote values written to the training env file (#6517)
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. Co-authored-by: Jaya Kasirajan <jkasiraj@amazon.com>
1 parent 94899a9 commit 8d07fea

8 files changed

Lines changed: 244 additions & 29 deletions

File tree

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import subprocess
2020
import json
2121
import os
22+
import shlex
2223
import sys
2324
from pathlib import Path
2425
import logging
@@ -238,7 +239,8 @@ def set_env(
238239
}
239240
with open(output_file, "w") as f:
240241
for key, value in env_vars.items():
241-
f.write(f"export {key}='{safe_serialize(value)}'\n")
242+
# The file is consumed with `source`, so the value must be shell-quoted.
243+
f.write(f"export {key}={shlex.quote(safe_serialize(value))}\n")
242244

243245
logger.info("Environment Variables:")
244246
log_env_variables(env_vars_dict=env_vars)

‎sagemaker-core/src/sagemaker/core/remote_function/runtime_environment/bootstrap_runtime_environment.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import multiprocessing
2121
import os
2222
import pathlib
23+
import shlex
2324
import shutil
2425
import subprocess
2526
import sys
@@ -527,7 +528,8 @@ def set_env(
527528

528529
with open(output_file, "w") as f:
529530
for key, value in env_vars.items():
530-
f.write(f"export {key}='{safe_serialize(value)}'\n")
531+
# The file is consumed with `source`, so the value must be shell-quoted.
532+
f.write(f"export {key}={shlex.quote(safe_serialize(value))}\n")
531533

532534
logger.info("Environment Variables:")
533535
log_env_variables(env_vars_dict=env_vars)

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,11 @@
1313

1414
import json
1515
import os
16+
import shutil
17+
import subprocess
18+
19+
import pytest
20+
1621
from unittest.mock import patch, mock_open
1722
from sagemaker.core.modules.train.container_drivers.scripts.environment import (
1823
num_cpus,
@@ -383,3 +388,61 @@ def test_set_env_multiple_channels(
383388

384389
# Verify file was written
385390
mock_file.assert_called_once()
391+
392+
393+
_CORE_ENV = "sagemaker.core.modules.train.container_drivers.scripts.environment"
394+
395+
_CORE_RESOURCE_CONFIG = {
396+
"current_host": "algo-1",
397+
"current_instance_type": "ml.m5.large",
398+
"hosts": ["algo-1"],
399+
"network_interface_name": "eth0",
400+
}
401+
402+
403+
def _core_write_env(hyperparameters, output_file):
404+
"""Run set_env with the given hyperparameters and no source code or distributed config."""
405+
patches = [
406+
patch(f"{_CORE_ENV}.read_source_code_json", return_value={}),
407+
patch(f"{_CORE_ENV}.read_distributed_json", return_value={}),
408+
patch(f"{_CORE_ENV}.num_cpus", return_value=1),
409+
patch(f"{_CORE_ENV}.num_gpus", return_value=0),
410+
patch(f"{_CORE_ENV}.num_neurons", return_value=0),
411+
]
412+
with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}):
413+
for p in patches:
414+
p.start()
415+
try:
416+
set_env(
417+
resource_config=_CORE_RESOURCE_CONFIG,
418+
input_data_config={},
419+
hyperparameters_config=hyperparameters,
420+
output_file=output_file,
421+
)
422+
finally:
423+
for p in patches:
424+
p.stop()
425+
426+
427+
@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file")
428+
@pytest.mark.parametrize(
429+
"payload",
430+
["it's 0.001'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"],
431+
)
432+
def test_set_env_hyperparameter_value_survives_source_verbatim(tmp_path, payload):
433+
"""A hyperparameter value must round-trip through `source` and never be executed."""
434+
marker = tmp_path / "marker"
435+
payload = payload.format(marker=marker)
436+
output_file = str(tmp_path / "sm_training.env")
437+
438+
_core_write_env({"learning_rate": payload}, output_file)
439+
440+
result = subprocess.run(
441+
["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_HP_LEARNING_RATE"'],
442+
capture_output=True,
443+
text=True,
444+
check=True,
445+
)
446+
447+
assert result.stdout == payload
448+
assert not marker.exists()

‎sagemaker-core/tests/unit/remote_function/runtime_environment/test_bootstrap_runtime_environment.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,10 @@
1111
# ANY KIND, either express or implied. See the License for the specific
1212
# language governing permissions and limitations under the License.
1313

14+
import os
15+
import shutil
16+
import subprocess
17+
1418
import pytest
1519
from unittest.mock import Mock, patch, mock_open
1620

@@ -543,3 +547,42 @@ def test_main_failure(self, mock_write_failure, mock_parse):
543547

544548
assert exc_info.value.code == DEFAULT_FAILURE_CODE
545549
mock_write_failure.assert_called_once()
550+
551+
552+
_BOOTSTRAP = "sagemaker.core.remote_function.runtime_environment.bootstrap_runtime_environment"
553+
554+
555+
@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file")
556+
@pytest.mark.parametrize(
557+
"payload",
558+
["eth0'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"],
559+
)
560+
def test_set_env_value_survives_source_verbatim(tmp_path, payload):
561+
"""A value written to the env file must round-trip through `source` and never be executed."""
562+
marker = tmp_path / "marker"
563+
payload = payload.format(marker=marker)
564+
output_file = str(tmp_path / "sm_training.env")
565+
resource_config = {
566+
"current_host": "algo-1",
567+
"current_instance_type": "ml.m5.large",
568+
"hosts": ["algo-1"],
569+
"network_interface_name": payload,
570+
}
571+
572+
with (
573+
patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}),
574+
patch(f"{_BOOTSTRAP}.num_cpus", return_value=1),
575+
patch(f"{_BOOTSTRAP}.num_gpus", return_value=0),
576+
patch(f"{_BOOTSTRAP}.num_neurons", return_value=0),
577+
):
578+
set_env(resource_config, output_file=output_file)
579+
580+
result = subprocess.run(
581+
["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_NETWORK_INTERFACE_NAME"'],
582+
capture_output=True,
583+
text=True,
584+
check=True,
585+
)
586+
587+
assert result.stdout == payload
588+
assert not marker.exists()

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import subprocess
2020
import json
2121
import os
22+
import shlex
2223
import sys
2324
from pathlib import Path
2425
import logging
@@ -238,7 +239,8 @@ def set_env(
238239
}
239240
with open(output_file, "w") as f:
240241
for key, value in env_vars.items():
241-
f.write(f"export {key}='{safe_serialize(value)}'\n")
242+
# The file is consumed with `source`, so the value must be shell-quoted.
243+
f.write(f"export {key}={shlex.quote(safe_serialize(value))}\n")
242244

243245
logger.info("Environment Variables:")
244246
log_env_variables(env_vars_dict=env_vars)

‎sagemaker-train/src/sagemaker/train/remote_function/runtime_environment/bootstrap_runtime_environment.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import multiprocessing
2121
import os
2222
import pathlib
23+
import shlex
2324
import shutil
2425
import subprocess
2526
import sys
@@ -524,7 +525,8 @@ def set_env(
524525

525526
with open(output_file, "w") as f:
526527
for key, value in env_vars.items():
527-
f.write(f"export {key}='{safe_serialize(value)}'\n")
528+
# The file is consumed with `source`, so the value must be shell-quoted.
529+
f.write(f"export {key}={shlex.quote(safe_serialize(value))}\n")
528530

529531
logger.info("Environment Variables:")
530532
log_env_variables(env_vars_dict=env_vars)

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

Lines changed: 85 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@
1717
import os
1818
import io
1919
import logging
20+
import shutil
21+
import subprocess
22+
23+
import pytest
2024

2125
from unittest.mock import patch
2226

@@ -87,37 +91,37 @@
8791

8892
# flake8: noqa
8993
EXPECTED_ENV = """
90-
export SM_MODEL_DIR='/opt/ml/model'
91-
export SM_INPUT_DIR='/opt/ml/input'
92-
export SM_INPUT_DATA_DIR='/opt/ml/input/data'
93-
export SM_INPUT_CONFIG_DIR='/opt/ml/input/config'
94-
export SM_OUTPUT_DIR='/opt/ml/output'
95-
export SM_OUTPUT_FAILURE='/opt/ml/output/failure'
96-
export SM_OUTPUT_DATA_DIR='/opt/ml/output/data'
97-
export SM_LOG_LEVEL='20'
98-
export SM_MASTER_ADDR='algo-1'
99-
export SM_MASTER_PORT='7777'
100-
export SM_SOURCE_DIR='/opt/ml/input/data/code'
101-
export SM_ENTRY_SCRIPT='train.py'
102-
export SM_DISTRIBUTED_DRIVER_DIR='/opt/ml/input/data/sm_drivers/distributed_drivers'
94+
export SM_MODEL_DIR=/opt/ml/model
95+
export SM_INPUT_DIR=/opt/ml/input
96+
export SM_INPUT_DATA_DIR=/opt/ml/input/data
97+
export SM_INPUT_CONFIG_DIR=/opt/ml/input/config
98+
export SM_OUTPUT_DIR=/opt/ml/output
99+
export SM_OUTPUT_FAILURE=/opt/ml/output/failure
100+
export SM_OUTPUT_DATA_DIR=/opt/ml/output/data
101+
export SM_LOG_LEVEL=20
102+
export SM_MASTER_ADDR=algo-1
103+
export SM_MASTER_PORT=7777
104+
export SM_SOURCE_DIR=/opt/ml/input/data/code
105+
export SM_ENTRY_SCRIPT=train.py
106+
export SM_DISTRIBUTED_DRIVER_DIR=/opt/ml/input/data/sm_drivers/distributed_drivers
103107
export SM_DISTRIBUTED_CONFIG='{"process_count_per_node": 2}'
104-
export SM_CHANNEL_TRAIN='/opt/ml/input/data/train'
105-
export SM_CHANNEL_VALIDATION='/opt/ml/input/data/validation'
108+
export SM_CHANNEL_TRAIN=/opt/ml/input/data/train
109+
export SM_CHANNEL_VALIDATION=/opt/ml/input/data/validation
106110
export SM_CHANNELS='["train", "validation"]'
107-
export SM_HP_BATCH_SIZE='32'
108-
export SM_HP_LEARNING_RATE='0.001'
111+
export SM_HP_BATCH_SIZE=32
112+
export SM_HP_LEARNING_RATE=0.001
109113
export SM_HP_HOSTS='["algo-1", "algo-2"]'
110114
export SM_HP_MP_PARAMETERS='{"microbatches": 2, "partitions": 2, "pipeline": "interleaved", "optimize": "memory", "horovod": true}'
111115
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}}'
112-
export SM_CURRENT_HOST='algo-1'
113-
export SM_CURRENT_INSTANCE_TYPE='ml.p3.16xlarge'
116+
export SM_CURRENT_HOST=algo-1
117+
export SM_CURRENT_INSTANCE_TYPE=ml.p3.16xlarge
114118
export SM_HOSTS='["algo-1", "algo-2", "algo-3"]'
115-
export SM_NETWORK_INTERFACE_NAME='eth0'
116-
export SM_HOST_COUNT='3'
117-
export SM_CURRENT_HOST_RANK='0'
118-
export SM_NUM_CPUS='8'
119-
export SM_NUM_GPUS='0'
120-
export SM_NUM_NEURONS='0'
119+
export SM_NETWORK_INTERFACE_NAME=eth0
120+
export SM_HOST_COUNT=3
121+
export SM_CURRENT_HOST_RANK=0
122+
export SM_NUM_CPUS=8
123+
export SM_NUM_GPUS=0
124+
export SM_NUM_NEURONS=0
121125
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"}'
122126
export SM_INPUT_DATA_CONFIG='{"train": {"ContentType": "trainingContentType", "TrainingInputMode": "File", "S3DistributionType": "FullyReplicated", "RecordWrapperType": "None"}, "validation": {"TrainingInputMode": "File", "S3DistributionType": "FullyReplicated", "RecordWrapperType": "None"}}'
123127
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():
212216
def _remove_extra_lines(string):
213217
"""Removes extra blank lines from a string."""
214218
return "\n".join([line for line in string.splitlines() if line.strip()])
219+
220+
221+
_ENV_PATCHES = [
222+
patch(
223+
"sagemaker.train.container_drivers.scripts.environment.read_source_code_json",
224+
return_value={},
225+
),
226+
patch(
227+
"sagemaker.train.container_drivers.scripts.environment.read_distributed_json",
228+
return_value={},
229+
),
230+
patch("sagemaker.train.container_drivers.scripts.environment.num_cpus", return_value=1),
231+
patch("sagemaker.train.container_drivers.scripts.environment.num_gpus", return_value=0),
232+
patch("sagemaker.train.container_drivers.scripts.environment.num_neurons", return_value=0),
233+
]
234+
235+
236+
def _write_env(hyperparameters, output_file):
237+
"""Run set_env with the given hyperparameters and no source code or distributed config."""
238+
with patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}):
239+
for p in _ENV_PATCHES:
240+
p.start()
241+
try:
242+
set_env(
243+
resource_config=RESOURCE_CONFIG,
244+
input_data_config=INPUT_DATA_CONFIG,
245+
hyperparameters_config=hyperparameters,
246+
output_file=output_file,
247+
)
248+
finally:
249+
for p in _ENV_PATCHES:
250+
p.stop()
251+
252+
253+
@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file")
254+
@pytest.mark.parametrize(
255+
"payload",
256+
["it's 0.001'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"],
257+
)
258+
def test_set_env_hyperparameter_value_survives_source_verbatim(tmp_path, payload):
259+
"""A hyperparameter value must round-trip through `source` and never be executed."""
260+
marker = tmp_path / "marker"
261+
payload = payload.format(marker=marker)
262+
output_file = str(tmp_path / "sm_training.env")
263+
264+
_write_env({"learning_rate": payload}, output_file)
265+
266+
result = subprocess.run(
267+
["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_HP_LEARNING_RATE"'],
268+
capture_output=True,
269+
text=True,
270+
check=True,
271+
)
272+
273+
assert result.stdout == payload
274+
assert not marker.exists()

‎sagemaker-train/tests/unit/train/remote_function/test_bootstrap_runtime_environment.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
from __future__ import absolute_import
1616

1717
import json
18+
import os
19+
import shutil
1820
import pytest
1921
import subprocess
2022
from unittest.mock import patch, MagicMock, mock_open
@@ -810,3 +812,42 @@ def test_main_non_root_user(self, mock_getuser, mock_manager_class):
810812
main(args)
811813

812814
mock_manager.change_dir_permission.assert_called_once()
815+
816+
817+
_BOOTSTRAP = "sagemaker.train.remote_function.runtime_environment.bootstrap_runtime_environment"
818+
819+
820+
@pytest.mark.skipif(shutil.which("bash") is None, reason="needs bash to source the env file")
821+
@pytest.mark.parametrize(
822+
"payload",
823+
["eth0'; touch {marker}; echo '", "'$(touch {marker})'", "x';touch {marker};#"],
824+
)
825+
def test_set_env_value_survives_source_verbatim(tmp_path, payload):
826+
"""A value written to the env file must round-trip through `source` and never be executed."""
827+
marker = tmp_path / "marker"
828+
payload = payload.format(marker=marker)
829+
output_file = str(tmp_path / "sm_training.env")
830+
resource_config = {
831+
"current_host": "algo-1",
832+
"current_instance_type": "ml.m5.large",
833+
"hosts": ["algo-1"],
834+
"network_interface_name": payload,
835+
}
836+
837+
with (
838+
patch.dict(os.environ, {"TRAINING_JOB_NAME": "test-job"}),
839+
patch(f"{_BOOTSTRAP}.num_cpus", return_value=1),
840+
patch(f"{_BOOTSTRAP}.num_gpus", return_value=0),
841+
patch(f"{_BOOTSTRAP}.num_neurons", return_value=0),
842+
):
843+
set_env(resource_config, output_file=output_file)
844+
845+
result = subprocess.run(
846+
["bash", "-c", f'source "{output_file}"; printf "%s" "$SM_NETWORK_INTERFACE_NAME"'],
847+
capture_output=True,
848+
text=True,
849+
check=True,
850+
)
851+
852+
assert result.stdout == payload
853+
assert not marker.exists()

0 commit comments

Comments
 (0)