feat: Add code_location support to Processing

This change gives users the ability to control upload location of
ScriptProcessor code with a code_location parameter, similarly to
Framework estimators and FrameworkProcessor. Extends both unit and
integration tests to verify correct upload locations.

Author: athewsey

src/sagemaker/processing.py

@@ -399,6 +399,7 @@ def __init__(
         env=None,
         tags=None,
         network_config=None,
+        code_location=None,
     ):
         """Initializes a ``ScriptProcessor`` instance.
 
@@ -443,10 +444,20 @@ def __init__(
                 A :class:`~sagemaker.network.NetworkConfig`
                 object that configures network isolation, encryption of
                 inter-container traffic, security group IDs, and subnets.
+            code_location (str): The S3 prefix URI where custom code will be uploaded
+                (default: None) - don't include a trailing slash since a string prepended
+                with a "/" is appended to ``code_location``. Code will be uploaded to
+                folder 'code_location/job-name/input/code'. If not specified, the default
+                ``code location`` is 's3://{sagemaker-default-bucket}'.
         """
         self._CODE_CONTAINER_BASE_PATH = "/opt/ml/processing/input/"
         self._CODE_CONTAINER_INPUT_NAME = "code"
         self.command = command
+        self.code_location = (
+            code_location[:-1]
+            if (code_location and code_location.endswith("/"))
+            else code_location
+        )
 
         super(ScriptProcessor, self).__init__(
             role=role,
@@ -653,9 +664,15 @@ def _upload_code(self, code, kms_key=None):
             str: The S3 URI of the uploaded file or directory.
 
         """
+        if self.code_location:
+            code_bucket, key_prefix = s3.parse_s3_url(self.code_location)
+        else:
+            code_bucket = self.sagemaker_session.default_bucket()
+            key_prefix = ""
         desired_s3_uri = s3.s3_path_join(
             "s3://",
-            self.sagemaker_session.default_bucket(),
+            code_bucket,
+            key_prefix,
             self._current_job_name,
             "input",
             self._CODE_CONTAINER_INPUT_NAME,
@@ -1373,15 +1390,12 @@ def __init__(
             env=env,
             tags=tags,
             network_config=network_config,
+            code_location=code_location,
         )
 
         self._FRAMEWORK_ENTRYPOINT_CONTAINER_INPUT_NAME = "entrypoint"
         self._FRAMEWORK_ENTRYPOINT_SCRIPT_NAME = "runproc.sh"
 
-        self.code_location = (
-            code_location[:-1] if (code_location and code_location.endswith("/")) else code_location
-        )
-
         if image_uri is None or base_job_name is None:
             # For these default configuration purposes, we don't need the optional args:
             est = self._create_estimator()

src/sagemaker/sklearn/processing.py

@@ -41,6 +41,7 @@ def __init__(
         env=None,
         tags=None,
         network_config=None,
+        code_location=None,
     ):
         """Initialize an ``SKLearnProcessor`` instance.
 
@@ -80,6 +81,11 @@ def __init__(
             network_config (sagemaker.network.NetworkConfig): A NetworkConfig
                 object that configures network isolation, encryption of
                 inter-container traffic, security group IDs, and subnets.
+            code_location (str): The S3 prefix URI where custom code will be uploaded
+                (default: None) - don't include a trailing slash since a string prepended
+                with a "/" is appended to ``code_location``. Code will be uploaded to
+                folder 'code_location/job-name/input/code'. If not specified, the default
+                ``code location`` is 's3://{sagemaker-default-bucket}'.
         """
         if not command:
             command = ["python3"]
@@ -106,4 +112,5 @@ def __init__(
             env=env,
             tags=tags,
             network_config=network_config,
+            code_location=code_location,
         )

tests/integ/test_processing.py

@@ -701,8 +701,9 @@ def test_sklearn_with_network_config(sagemaker_session, sklearn_latest_version,
 
 
 def test_processing_job_inputs_and_output_config(
-    sagemaker_session, image_uri, cpu_instance_type, output_kms_key
+    sagemaker_session, image_uri, cpu_instance_type, output_kms_key, custom_bucket_name,
 ):
+    custom_code_location = f"s3://{custom_bucket_name}/customized-processing-code"
     script_processor = ScriptProcessor(
         role=ROLE,
         image_uri=image_uri,
@@ -716,6 +717,7 @@ def test_processing_job_inputs_and_output_config(
         base_job_name="test-script-processor",
         env={"DUMMY_ENVIRONMENT_VARIABLE": "dummy-value"},
         tags=[{"Key": "dummy-tag", "Value": "dummy-tag-value"}],
+        code_location=custom_code_location,
         sagemaker_session=sagemaker_session,
     )
 
@@ -734,6 +736,10 @@ def test_processing_job_inputs_and_output_config(
     assert (
         job_description["ProcessingInputs"][:-1] == expected_inputs_and_outputs["ProcessingInputs"]
     )
+    assert job_description["ProcessingInputs"][-1]["InputName"] == "code"
+    assert job_description["ProcessingInputs"][-1]["S3Input"]["S3Uri"].startswith(
+        custom_code_location
+    )
     assert (
         job_description["ProcessingOutputConfig"]
         == expected_inputs_and_outputs["ProcessingOutputConfig"]

tests/integ/test_s3.py

@@ -17,8 +17,7 @@
 
 import pytest
 
-from sagemaker.s3 import S3Uploader
-from sagemaker.s3 import S3Downloader
+from sagemaker.s3 import S3Downloader, S3Uploader, s3_path_join
 
 from tests.integ.kms_utils import get_or_create_kms_key
 
@@ -238,3 +237,8 @@ def test_s3_uploader_and_downloader_downloads_files_when_given_directory_uris_wi
 
     with open(os.path.join(TMP_BASE_PATH, my_inner_directory_uuid, file_2_name), "r") as f:
         assert file_2_body == f.read()
+
+
+def test_s3_path_join_ignores_empty_elements():
+    # (At writing, ScriptProcessor code upload expects/requires this)
+    assert s3_path_join("s3://", "mybucket", "", "a", "b", "", "c") == "s3://mybucket/a/b/c"