Merge pull request aws#6 from athewsey/feat/fw-processor

Restore 'command' for FrameworkProcessors

Author: verdimrc

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