From ab05ab3f5f28d67c254ff849ce09c54d9a2156e7 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 28 Sep 2026 14:26:24 -0700 Subject: [PATCH 1/2] feat: add args field to SourceCode for command-based training (#5226) --- .../src/sagemaker/core/modules/configs.py | 8 + .../src/sagemaker/core/training/configs.py | 8 + .../src/sagemaker/train/model_trainer.py | 17 +- .../tests/unit/train/test_sourcecode_args.py | 178 ++++++++++++++++++ 4 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 sagemaker-train/tests/unit/train/test_sourcecode_args.py 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/tests/unit/train/test_sourcecode_args.py b/sagemaker-train/tests/unit/train/test_sourcecode_args.py new file mode 100644 index 0000000000..0066e308ff --- /dev/null +++ b/sagemaker-train/tests/unit/train/test_sourcecode_args.py @@ -0,0 +1,178 @@ +# 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 os +import logging +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() + + +def test_sourcecode_args_field_defaults_to_none(): + """The new ``args`` field exists and defaults to None (backward compatible).""" + assert SourceCode(command="python train.py").args is None + + +def test_sourcecode_args_field_synced_in_modules_configs(): + """``args`` must exist in the legacy ``modules.configs`` mirror as well.""" + from sagemaker.core.modules.configs import SourceCode as ModulesSourceCode + + assert ModulesSourceCode(command="python train.py").args is None + assert ModulesSourceCode(command="python train.py", args=["--epochs", "5"]).args == [ + "--epochs", + "5", + ] + + +def test_sourcecode_args_accepts_numeric_values(): + """``args`` accepts int/float values (matching the maintainer's proposed API).""" + # Both SourceCode definitions must accept the issue-thread example without raising. + from sagemaker.core.modules.configs import SourceCode as ModulesSourceCode + + for cls in (SourceCode, ModulesSourceCode): + source_code = cls(command="python train.py", args=["--epochs", 25, "--lr", 0.001]) + assert source_code.args == ["--epochs", 25, "--lr", 0.001] + + +def test_prepare_train_script_appends_args_to_command(): + """``args`` are appended to ``command`` in the executed base command.""" + source_code = SourceCode( + source_dir="scripts", + command="python launcher.py -e train.py", + args=["--epochs", "25", "--learning_rate", "0.001"], + ) + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + assert 'CMD="python launcher.py -e train.py --epochs 25 --learning_rate 0.001"' in script + + +def test_prepare_train_script_appends_numeric_args(): + """Numeric ``args`` are stringified and appended to the executed command.""" + source_code = SourceCode( + source_dir="scripts", + command="python launcher.py", + args=["--epochs", 25, "--learning_rate", 0.001], + ) + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + assert 'CMD="python launcher.py --epochs 25 --learning_rate 0.001"' in script + + +def test_prepare_train_script_shell_quotes_args_with_spaces(): + """Args with spaces/special chars are shell-quoted so they stay single arguments.""" + source_code = SourceCode( + source_dir="scripts", + command="python launcher.py", + args=["--prompt", "hello world"], + ) + trainer = _make_trainer(source_code) + script = _written_train_script(trainer, source_code) + # shlex.quote wraps the value so the space does not re-tokenize the argument. + assert "'hello world'" in script + assert "CMD=\"python launcher.py --prompt 'hello world'\"" in script + + +def test_prepare_train_script_command_only_no_regression(): + """No ``args`` -> 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 'CMD="python launcher.py -e train.py"' in script + + +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) From 04a31454030a7f267c3ecc3ad110abf61123322b Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Wed, 30 Sep 2026 08:50:52 -0700 Subject: [PATCH 2/2] fix: stop the generated train script from defeating shlex.quote on args The `args` shell-quoting introduced with SourceCode.args did not actually hold, because EXECUTE_BASE_COMMANDS interpolated the already-quoted command inside a double-quoted assignment: CMD="{base_command}" eval $CMD Single quotes are inert inside double quotes, so bash expanded the contents at assignment time and the quoting was gone before `eval` ever ran. Verified against real bash, an arg of: $HOME arrived as /Users/... (expected literal $HOME) $(id -u) arrived as 503 (expected literal $(id -u)) `id -u` arrived as 503 (expected literal `id -u`) a"b aborted the script with `syntax error: unexpected end of file` it's aborted the script likewise a\tb, a\nb collapsed to a b So the docstring promise that values "may contain spaces or special characters" held only for whitespace, which is the one case the original tests covered -- that is why all nine of them passed against a broken implementation. The command is now carried through a quoted heredoc, which performs no expansion whatsoever, and `eval "$CMD"` parses it exactly once: CMD=$(cat <<'SAGEMAKER_BASE_COMMAND_EOF' {base_command} SAGEMAKER_BASE_COMMAND_EOF ) eval "$CMD" All 15 metacharacter cases now round-trip verbatim. `eval` is deliberately retained so `command` keeps its existing semantics: env vars in the command still expand and shell operators like `&&` still chain. Only the spurious *second* expansion -- the one that defeated the quoting -- is removed. Tests now assert on behavior instead of on the text of the generated script: they execute the generated command block under bash against a stub that reports its argv, which is the only way to pin a quoting contract. Two graded negative controls confirm they do the work: - reverting to CMD="..." -> 23 of 28 fail - keeping the heredoc but unquoting the delimiter (expansion only) -> exactly the 4 expansion cases fail The five assertions that string-matched `CMD="..."` are rewritten against a helper that tolerates either delimiter form, so a future regression fails on semantics rather than on a parse error in the test. No regressions: 2282 passed in the sagemaker-train unit suite. The unused copy in sagemaker-core/src/sagemaker/core/modules/templates.py has the same flaw but is imported by nothing, so it is left out of scope here. --- .../src/sagemaker/train/templates.py | 13 +- .../tests/unit/train/test_sourcecode_args.py | 182 +++++++++++++++++- 2 files changed, 188 insertions(+), 7 deletions(-) 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 index 0066e308ff..a42ad9afae 100644 --- a/sagemaker-train/tests/unit/train/test_sourcecode_args.py +++ b/sagemaker-train/tests/unit/train/test_sourcecode_args.py @@ -21,8 +21,12 @@ from __future__ import absolute_import +import json import os import logging +import re +import shutil +import subprocess from tempfile import TemporaryDirectory import pytest @@ -88,6 +92,68 @@ def _written_train_script(trainer, source_code): 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 <