fix: restore 'command' for FrameworkProcessors

FrameworkProcessor now accepts a 'command' param like the
ScriptProcessor class did previously, and tries to maintain similar
API by forwarding it to the parent class... But patching the
container run command generation from the parent to override. Also
add initial unit tests for multi-instance and overridden commands.

Author: athewsey

src/sagemaker/mxnet/processing.py

@@ -34,6 +34,7 @@ def __init__(
         instance_type,
         py_version="py3",  # New kwarg
         image_uri=None,
+        command=["python"],
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
@@ -66,6 +67,7 @@ def __init__(
             instance_type,
             py_version,
             image_uri,
+            command,
             volume_size_in_gb,
             volume_kms_key,
             output_kms_key,

src/sagemaker/processing.py

@@ -21,6 +21,7 @@
 import os
 import pathlib
 import logging
+from textwrap import dedent
 from typing import Dict, List, Optional, Tuple
 import attr
 
@@ -1217,18 +1218,7 @@ class FeatureStoreOutput(ApiObject):
 class FrameworkProcessor(ScriptProcessor):
     """Handles Amazon SageMaker processing tasks for jobs using a machine learning framework."""
 
-    runproc_sh = """#!/bin/bash
-
-cd /opt/ml/processing/input/code/
-tar -xzf sourcedir.tar.gz
-
-# Exit on any error. SageMaker uses error code to mark failed job.
-set -e
-
-[[ -f 'requirements.txt' ]] && pip install -r requirements.txt
-
-python {entry_point} "$@"
-"""
+    framework_entrypoint_command = ["/bin/bash"]
 
     # Added new (kw)args for estimator. The rest are from ScriptProcessor with same defaults.
     def __init__(
@@ -1240,6 +1230,7 @@ def __init__(
         instance_type,
         py_version="py3",  # New kwarg
         image_uri=None,
+        command=["python"],
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
@@ -1272,6 +1263,8 @@ def __init__(
                 is ignored when ``image_uri`` is provided.
             image_uri (str): The URI of the Docker image to use for the
                 processing jobs (default: None).
+            command ([str]): The command to run, along with any command-line flags
+                to *precede* the ```entry_point script``` (default: ['python']).
             volume_size_in_gb (int): Size in GB of the EBS volume
                 to use for storing data during processing (default: 30).
             volume_kms_key (str): A KMS key for the processing volume (default: None).
@@ -1312,7 +1305,7 @@ def __init__(
         super().__init__(
             role=role,
             image_uri=image_uri,
-            command=["/bin/bash"],
+            command=command,
             instance_count=instance_count,
             instance_type=instance_type,
             volume_size_in_gb=volume_size_in_gb,
@@ -1493,7 +1486,7 @@ def run(  # type: ignore[override]
         )
         script = estimator.uploaded_code.script_name
         s3_runproc_sh = S3Uploader.upload_string_as_file_body(
-            self.runproc_sh.format(entry_point=script),
+            self._generate_framework_script(script),
             desired_s3_uri=entrypoint_s3_uri,
             sagemaker_session=self.sagemaker_session,
         )
@@ -1512,6 +1505,35 @@ def run(  # type: ignore[override]
             kms_key=kms_key,
         )
 
+    def _generate_framework_script(self, user_script: str) -> str:
+        """Generate the framework entrypoint file (as text) for a processing job.
+
+        This script implements the "framework" functionality for setting up your code:
+        Untar-ing the sourcedir bundle in the ```code``` input; installing extra
+        runtime dependencies if specified; and then invoking the ```command``` and
+        ```entry_point``` configured for the job.
+
+        Args:
+            user_script (str): Relative path to ```entry_point``` in the source bundle
+                - e.g. 'process.py'.
+        """
+        return dedent("""\
+            #!/bin/bash
+
+            cd /opt/ml/processing/input/code/
+            tar -xzf sourcedir.tar.gz
+
+            # Exit on any error. SageMaker uses error code to mark failed job.
+            set -e
+
+            [[ -f 'requirements.txt' ]] && pip install -r requirements.txt
+
+            {entry_point_command} {entry_point} "$@"
+        """).format(
+            entry_point_command=" ".join(self.command),
+            entry_point=user_script,
+        )
+
     def _upload_payload(
         self,
         entry_point: str,
@@ -1575,3 +1597,18 @@ def _patch_inputs_with_payload(self, inputs, s3_payload) -> List[ProcessingInput
             )
         )
         return inputs
+
+    def _set_entrypoint(self, command, user_script_name):
+        """FrameworkProcessor override for setting processing job entrypoint.
+
+        Args:
+            command ([str]): Ignored in favor of self.framework_entrypoint_command
+            user_script_name (str): A filename with an extension.
+        """
+        
+        user_script_location = str(
+            pathlib.PurePosixPath(
+                self._CODE_CONTAINER_BASE_PATH, self._CODE_CONTAINER_INPUT_NAME, user_script_name
+            )
+        )
+        self.entrypoint = self.framework_entrypoint_command + [user_script_location]

src/sagemaker/pytorch/processing.py

@@ -34,6 +34,7 @@ def __init__(
         instance_type,
         py_version="py3",  # New kwarg
         image_uri=None,
+        command=["python"],
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
@@ -66,6 +67,7 @@ def __init__(
             instance_type,
             py_version,
             image_uri,
+            command,
             volume_size_in_gb,
             volume_kms_key,
             output_kms_key,

src/sagemaker/sklearn/processing.py

@@ -48,6 +48,7 @@ def __init__(
         instance_type,
         py_version="py3",  # New kwarg
         image_uri=None,
+        command=["python"],
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
@@ -68,6 +69,7 @@ def __init__(
             instance_type,
             py_version,
             image_uri,
+            command,
             volume_size_in_gb,
             volume_kms_key,
             output_kms_key,

src/sagemaker/tensorflow/processing.py

@@ -34,6 +34,7 @@ def __init__(
         instance_type,
         py_version="py3",  # New kwarg
         image_uri=None,
+        command=["python"],
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
@@ -66,6 +67,7 @@ def __init__(
             instance_type,
             py_version,
             image_uri,
+            command,
             volume_size_in_gb,
             volume_kms_key,
             output_kms_key,

src/sagemaker/xgboost/processing.py

@@ -34,6 +34,7 @@ def __init__(
         instance_type,
         py_version="py3",  # New kwarg
         image_uri=None,
+        command=["python"],
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
@@ -66,6 +67,7 @@ def __init__(
             instance_type,
             py_version,
             image_uri,
+            command,
             volume_size_in_gb,
             volume_kms_key,
             output_kms_key,

tests/unit/test_processing.py

@@ -106,6 +106,7 @@ def test_sklearn_with_all_parameters(
     processor = SKLearnProcessor(
         role=ROLE,
         framework_version=sklearn_version,
+        command=["Rscript"],
         instance_type="ml.m4.xlarge",
         instance_count=1,
         volume_size_in_gb=100,
@@ -153,12 +154,14 @@ def test_sklearn_with_all_parameters_via_run_args(
     exists_mock, isfile_mock, botocore_resolver, sklearn_version, sagemaker_session
 ):
     botocore_resolver.return_value.construct_endpoint.return_value = {"hostname": ECR_HOSTNAME}
+    custom_command = ["Rscript"]
 
     processor = SKLearnProcessor(
         role=ROLE,
         framework_version=sklearn_version,
+        command=custom_command,
         instance_type="ml.m4.xlarge",
-        instance_count=1,
+        instance_count=2,
         volume_size_in_gb=100,
         volume_kms_key="arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
         output_kms_key="arn:aws:kms:us-west-2:012345678901:key/output-kms-key",
@@ -195,14 +198,27 @@ def test_sklearn_with_all_parameters_via_run_args(
         experiment_config={"ExperimentName": "AnExperiment"},
     )
 
-    expected_args = _get_expected_args_all_parameters_modular_code(processor._current_job_name)
+    expected_args = _get_expected_args_all_parameters_modular_code(
+        processor._current_job_name,
+        instance_count=2,
+    )
     sklearn_image_uri = (
         "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:{}-cpu-py3"
     ).format(sklearn_version)
     expected_args["app_specification"]["ImageUri"] = sklearn_image_uri
 
     sagemaker_session.process.assert_called_with(**expected_args)
 
+    # Verify the alternate command was applied successfully:
+    framework_script = processor._generate_framework_script("processing_code.py")
+    expected_invocation = f"{' '.join(custom_command)} processing_code.py"
+    assert f"\n{expected_invocation}" in framework_script, (
+        "Framework script should contain customized invocation:\n{}\n\nGot:\n{}".format(
+            expected_invocation,
+            framework_script,
+        )
+    )
+
 
 @patch("sagemaker.utils._botocore_resolver")
 @patch("os.path.exists", return_value=True)
@@ -811,7 +827,7 @@ def _get_data_outputs_all_parameters():
     ]
 
 
-def _get_expected_args_all_parameters_modular_code(job_name, code_s3_uri=MOCKED_S3_URI):
+def _get_expected_args_all_parameters_modular_code(job_name, code_s3_uri=MOCKED_S3_URI, instance_count=1):
     # Add something to inputs
     return {
         "inputs": [
@@ -927,7 +943,7 @@ def _get_expected_args_all_parameters_modular_code(job_name, code_s3_uri=MOCKED_
         "resources": {
             "ClusterConfig": {
                 "InstanceType": "ml.m4.xlarge",
-                "InstanceCount": 1,
+                "InstanceCount": instance_count,
                 "VolumeSizeInGB": 100,
                 "VolumeKmsKeyId": "arn:aws:kms:us-west-2:012345678901:key/volume-kms-key",
             }