diff --git a/sagemaker-core/src/sagemaker/core/modules/configs.py b/sagemaker-core/src/sagemaker/core/modules/configs.py index 865018f50c..90381f49ca 100644 --- a/sagemaker-core/src/sagemaker/core/modules/configs.py +++ b/sagemaker-core/src/sagemaker/core/modules/configs.py @@ -104,6 +104,13 @@ class SourceCode(BaseConfig): command (Optional[str]): The command(s) to execute in the training job container. Example: "python my_script.py". If not specified, entry_script must be provided. + args (Optional[List[Union[str, int, float]]]): + A list of arguments to append to ``command`` when it is executed in the training job + container. Example: ``["--epochs", 25, "--learning_rate", 0.001]``. Each argument is + shell-quoted before being appended, so values may contain spaces or special characters. + Only applicable when ``command`` is provided; top-level ``hyperparameters`` are not + passed as CLI arguments in ``command`` mode -- they are available inside the container + via the ``SM_HPS`` environment variable. ignore_patterns: (Optional[List[str]]) : The ignore patterns to ignore specific files/folders when uploading to S3. If not specified, default to: ['.env', '.git', '__pycache__', '.DS_Store', '.cache', '.ipynb_checkpoints']. @@ -113,6 +120,7 @@ class SourceCode(BaseConfig): requirements: Optional[str] = None entry_script: Optional[str] = None command: Optional[str] = None + args: Optional[List[Union[str, int, float]]] = None ignore_patterns: Optional[List[str]] = [ ".env", ".git", diff --git a/sagemaker-core/src/sagemaker/core/training/configs.py b/sagemaker-core/src/sagemaker/core/training/configs.py index 17e40b0771..46759a82c2 100644 --- a/sagemaker-core/src/sagemaker/core/training/configs.py +++ b/sagemaker-core/src/sagemaker/core/training/configs.py @@ -107,6 +107,13 @@ class SourceCode(BaseConfig): command (Optional[StrPipeVar]): The command(s) to execute in the training job container. Example: "python my_script.py". If not specified, entry_script must be provided. + args (Optional[List[Union[str, int, float]]]): + A list of arguments to append to ``command`` when it is executed in the training job + container. Example: ``["--epochs", 25, "--learning_rate", 0.001]``. Each argument is + shell-quoted before being appended, so values may contain spaces or special characters. + Only applicable when ``command`` is provided; top-level ``hyperparameters`` are not + passed as CLI arguments in ``command`` mode -- they are available inside the container + via the ``SM_HPS`` environment variable. ignore_patterns: (Optional[List[str]]) : The ignore patterns to ignore specific files/folders when uploading to S3. If not specified, default to: ['.env', '.git', '__pycache__', '.DS_Store', '.cache', '.ipynb_checkpoints']. @@ -116,6 +123,7 @@ class SourceCode(BaseConfig): requirements: Optional[StrPipeVar] = None entry_script: Optional[StrPipeVar] = None command: Optional[StrPipeVar] = None + args: Optional[List[Union[str, int, float]]] = None ignore_patterns: Optional[List[str]] = [ ".env", ".git", diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 35b1395cce..53a5958773 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -18,6 +18,7 @@ import os import json import re +import shlex import shutil from tempfile import TemporaryDirectory from typing import Optional, List, Union, Dict, Any, ClassVar @@ -1269,8 +1270,20 @@ def _prepare_train_script( "Both 'command' and 'entry_script' are provided in the SourceCode. " "Defaulting to 'command'." ) - base_command = source_code.command.split() - base_command = " ".join(base_command) + command_parts = source_code.command.split() + if source_code.args: + # Shell-quote each arg so values with spaces or special characters are + # passed through as single, literal arguments to the executed command. + command_parts.extend(shlex.quote(str(arg)) for arg in source_code.args) + base_command = " ".join(command_parts) + if self.hyperparameters: + logger.warning( + "Hyperparameters are set but are not passed as command-line arguments " + "when 'command' is used in the SourceCode. They are available inside the " + "training container via the 'SM_HPS' environment variable, e.g. " + "`hps = json.loads(os.environ['SM_HPS'])`. To pass arguments to your " + "'command', use the SourceCode 'args' field." + ) install_requirements = "" if source_code.requirements: diff --git a/sagemaker-train/src/sagemaker/train/templates.py b/sagemaker-train/src/sagemaker/train/templates.py index 2ae3ae766f..1aaae150ec 100644 --- a/sagemaker-train/src/sagemaker/train/templates.py +++ b/sagemaker-train/src/sagemaker/train/templates.py @@ -14,10 +14,19 @@ from __future__ import absolute_import +# The command is carried in through a quoted heredoc ('SAGEMAKER_BASE_COMMAND_EOF'), which +# performs no expansion at all, so the shell quoting applied by +# ModelTrainer._prepare_train_script survives into CMD verbatim. Assigning with +# CMD="{base_command}" instead would expand $VAR, $(...) and backticks at assignment time and +# let an embedded double quote terminate the string early, which defeats that quoting and +# corrupts SourceCode.args. `eval "$CMD"` then parses the command exactly once. EXECUTE_BASE_COMMANDS = """ -CMD="{base_command}" +CMD=$(cat <<'SAGEMAKER_BASE_COMMAND_EOF' +{base_command} +SAGEMAKER_BASE_COMMAND_EOF +) echo "Executing command: $CMD" -eval $CMD +eval "$CMD" """ EXECUTE_BASIC_SCRIPT_DRIVER = """ diff --git a/sagemaker-train/tests/unit/train/test_sourcecode_args.py b/sagemaker-train/tests/unit/train/test_sourcecode_args.py new file mode 100644 index 0000000000..a42ad9afae --- /dev/null +++ b/sagemaker-train/tests/unit/train/test_sourcecode_args.py @@ -0,0 +1,350 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Tests for the SourceCode ``args`` field (issue #5226). + +When ``command`` is used in ``SourceCode`` the SDK runs it verbatim and does NOT +append top-level ``hyperparameters`` as CLI arguments. ``args`` lets the user pass +positional/optional arguments that are appended to ``command`` deterministically, +and a warning points users to the ``SM_HPS`` env var when they combine +``hyperparameters`` with ``command``/``args``. +""" + +from __future__ import absolute_import + +import json +import os +import logging +import re +import shutil +import subprocess +from tempfile import TemporaryDirectory + +import pytest +from unittest.mock import patch, MagicMock + +from sagemaker.core.helper.session_helper import Session +from sagemaker.train.model_trainer import ModelTrainer +from sagemaker.train.configs import Compute, StoppingCondition, OutputDataConfig, SourceCode +from sagemaker.train.constants import TRAIN_SCRIPT + +DEFAULT_BASE_NAME = "dummy-image-job" +DEFAULT_IMAGE = "000000000000.dkr.ecr.us-west-2.amazonaws.com/dummy-image:latest" +DEFAULT_BUCKET = "sagemaker-us-west-2-000000000000" +DEFAULT_ROLE = "arn:aws:iam::000000000000:role/test-role" +DEFAULT_BUCKET_PREFIX = "sample-prefix" +DEFAULT_REGION = "us-west-2" +DEFAULT_COMPUTE_CONFIG = Compute(instance_type="ml.m5.xlarge", instance_count=1) +DEFAULT_OUTPUT_DATA_CONFIG = OutputDataConfig( + s3_output_path=f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}/{DEFAULT_BASE_NAME}", + compression_type="GZIP", + kms_key_id=None, +) +DEFAULT_STOPPING_CONDITION = StoppingCondition( + max_runtime_in_seconds=3600, + max_pending_time_in_seconds=None, + max_wait_time_in_seconds=None, +) + + +@pytest.fixture(autouse=True) +def modules_session(): + with ( + patch("sagemaker.train.Session", spec=Session) as session_mock, + patch("sagemaker.train.defaults.resolve_and_validate_role", return_value=DEFAULT_ROLE), + ): + session_instance = session_mock.return_value + session_instance.default_bucket.return_value = DEFAULT_BUCKET + session_instance.get_caller_identity_arn.return_value = DEFAULT_ROLE + session_instance.default_bucket_prefix = DEFAULT_BUCKET_PREFIX + session_instance.boto_session = MagicMock(spec="boto3.session.Session") + session_instance.boto_region_name = DEFAULT_REGION + yield session_instance + + +def _make_trainer(source_code, hyperparameters=None): + return ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + compute=DEFAULT_COMPUTE_CONFIG, + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + source_code=source_code, + hyperparameters=hyperparameters or {}, + ) + + +def _written_train_script(trainer, source_code): + with TemporaryDirectory() as tmp_name: + tmp_dir = MagicMock() + tmp_dir.name = tmp_name + trainer._prepare_train_script(tmp_dir=tmp_dir, source_code=source_code) + with open(os.path.join(tmp_name, TRAIN_SCRIPT), "r") as f: + return f.read() + + +_EVAL_LINE = 'eval "$CMD"' +# Deliberately tolerant of how the heredoc delimiter is written, so that a regression which +# keeps the structure but drops the quotes around the delimiter (re-enabling expansion) is +# caught by the semantic assertions below rather than by a parse error in this helper. +_HEREDOC_OPEN_RE = re.compile(r"^CMD=\$\(cat <<'?(?P\w+)'?$") + + +def _command_block_bounds(lines): + """Return (heredoc-open index, body slice, index just past ``eval``).""" + for index, line in enumerate(lines): + match = _HEREDOC_OPEN_RE.match(line) + if match: + delimiter = match.group("delimiter") + close = lines.index(delimiter, index + 1) + end = lines.index(_EVAL_LINE, close) + 1 + return index, slice(index + 1, close), end + raise AssertionError( + "no 'CMD=$(cat < command runs verbatim, unchanged from prior behavior.""" + source_code = SourceCode(source_dir="scripts", command="python launcher.py -e train.py") + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + assert _base_command(script) == "python launcher.py -e train.py" + + +requires_bash = pytest.mark.skipif( + shutil.which("bash") is None, reason="needs bash to execute the generated command block" +) + + +@requires_bash +@pytest.mark.parametrize( + "value", + [ + "hello world", + "$HOME", + "${HOME}", + "$(id -u)", + "`id -u`", + 'a"b', + "it's", + "; echo INJECTED", + "&& echo INJECTED", + "| cat", + "*", + "a\tb", + "a\nb", + "back\\slash", + "", + ], +) +def test_args_reach_the_command_as_literal_single_arguments(value): + """Each arg must arrive verbatim as one argument, whatever shell metacharacters it holds. + + Regression test for the quoting bug in the original implementation: ``shlex.quote`` was + defeated because the generated script assigned the command inside a double-quoted + ``CMD="..."``, so the shell expanded ``$VAR``/``$(...)``/backticks at assignment time and + an embedded double quote terminated the string early. Asserting on the text of the script + could not catch that -- only running it can. + """ + source_code = SourceCode(source_dir="scripts", command="capture_args", args=["--prompt", value]) + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + + assert _argv_from_running_base_command(script) == ["--prompt", value] + + +@requires_bash +def test_multiple_args_preserve_order_and_boundaries(): + """Adjacent args stay distinct even when each contains spaces.""" + values = ["--a", "one two", "--b", "three four", "--c", "5"] + source_code = SourceCode(source_dir="scripts", command="capture_args", args=values) + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + + assert _argv_from_running_base_command(script) == values + + +@requires_bash +def test_numeric_args_reach_the_command_as_strings(): + """int/float args are stringified without gaining quotes or losing precision.""" + source_code = SourceCode( + source_dir="scripts", + command="capture_args", + args=["--epochs", 25, "--learning_rate", 0.001], + ) + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + + assert _argv_from_running_base_command(script) == ["--epochs", "25", "--learning_rate", "0.001"] + + +@requires_bash +def test_environment_variables_in_command_still_expand(): + """``command`` itself is still shell-evaluated, so env vars in it keep working. + + The heredoc stops expansion at *assignment* time only; ``eval`` still parses the command + once, which is the behavior users of ``command`` rely on. + """ + source_code = SourceCode(source_dir="scripts", command="capture_args $HOME") + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + + argv = _argv_from_running_base_command(script) + assert argv == [os.environ["HOME"]] + + +@requires_bash +def test_shell_operators_in_command_still_work(): + """``&&`` in a user-supplied ``command`` still chains, unchanged from prior behavior.""" + source_code = SourceCode(source_dir="scripts", command="capture_args one && capture_args two") + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + + lines = script.splitlines() + start, _, end = _command_block_bounds(lines) + completed = subprocess.run( + ["bash", "-c", _CAPTURE_ARGV_PRELUDE + "\n".join(lines[start:end])], + capture_output=True, + text=True, + check=True, + ) + emitted = [ + json.loads(line[len("ARGV=") :]) + for line in completed.stdout.splitlines() + if line.startswith("ARGV=") + ] + assert emitted == [["one"], ["two"]] + + +def test_warning_when_hyperparameters_combined_with_command(caplog): + """A warning pointing to SM_HPS fires when hyperparameters + command are combined.""" + source_code = SourceCode(source_dir="scripts", command="python launcher.py -e train.py") + trainer = _make_trainer(source_code, hyperparameters={"epochs": 25}) + with caplog.at_level(logging.WARNING): + _written_train_script(trainer, source_code) + assert any("SM_HPS" in rec.message for rec in caplog.records) + + +def test_no_warning_when_command_without_hyperparameters(caplog): + """No SM_HPS warning when command is used but no hyperparameters are set.""" + source_code = SourceCode(source_dir="scripts", command="python launcher.py -e train.py") + trainer = _make_trainer(source_code) + with caplog.at_level(logging.WARNING): + _written_train_script(trainer, source_code) + assert not any("SM_HPS" in rec.message for rec in caplog.records)