feat: add args field to SourceCode for command-based training (#5226)

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",

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",

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:

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 = """