fix: allow processing users to run code in s3 (#1167)

* fix: allow processing users to run code in s3

Add validation and unit tests for Processing

* Update src/sagemaker/processing.py

Co-Authored-By: Karim Nakad <[email protected]>

Author: 2 people

src/sagemaker/processing.py

@@ -369,8 +369,9 @@ def run(
         """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
 
-        user_script_name = self._get_user_script_name(code)
-        user_code_s3_uri = self._upload_code(code)
+        user_code_s3_uri = self._handle_user_code_url(code)
+        user_script_name = self._get_user_code_name(code)
+
         inputs_with_code = self._convert_code_and_add_to_inputs(inputs, user_code_s3_uri)
 
         self._set_entrypoint(self.command, user_script_name)
@@ -389,25 +390,59 @@ def run(
         if wait:
             self.latest_job.wait(logs=logs)
 
-    def _get_user_script_name(self, code):
-        """Finds the user script name using the provided code file,
-        directory, or script name.
+    def _get_user_code_name(self, code):
+        """Gets the basename of the user's code from the URL the customer provided.
 
         Args:
-            code (str): This can be an S3 uri or a local path to either
-                a directory or a file.
+            code (str): A URL to the user's code.
+
+        Returns:
+            str: The basename of the user's code.
+
+        """
+        code_url = urlparse(code)
+        return os.path.basename(code_url.path)
+
+    def _handle_user_code_url(self, code):
+        """Gets the S3 URL containing the user's code.
+
+           Inspects the scheme the customer passed in ("s3://" for code in S3, "file://" or nothing
+           for absolute or local file paths. Uploads the code to S3 if the code is a local file.
+
+        Args:
+            code (str): A URL to the customer's code.
 
         Returns:
-            str: The script name from the S3 uri or from the file found
-                on the user's local machine.
+            str: The S3 URL to the customer's code.
+
         """
-        if os.path.isdir(code) is None or not os.path.splitext(code)[1]:
+        code_url = urlparse(code)
+        if code_url.scheme == "s3":
+            user_code_s3_uri = code
+        elif code_url.scheme == "" or code_url.scheme == "file":
+            # Validate that the file exists locally and is not a directory.
+            if not os.path.exists(code):
+                raise ValueError(
+                    """code {} wasn't found. Please make sure that the file exists.
+                    """.format(
+                        code
+                    )
+                )
+            if not os.path.isfile(code):
+                raise ValueError(
+                    """code {} must be a file, not a directory. Please pass a path to a file.
+                    """.format(
+                        code
+                    )
+                )
+            user_code_s3_uri = self._upload_code(code)
+        else:
             raise ValueError(
-                """'code' must be a file, not a directory. Please pass a path to a file, not a
-                directory.
-                """
+                "code {} url scheme {} is not recognized. Please pass a file path or S3 url".format(
+                    code, code_url.scheme
+                )
             )
-        return os.path.basename(code)
+        return user_code_s3_uri
 
     def _upload_code(self, code):
         """Uploads a code file or directory specified as a string

src/sagemaker/sklearn/processing.py

@@ -26,6 +26,8 @@
 class SKLearnProcessor(ScriptProcessor):
     """Handles Amazon SageMaker processing tasks for jobs using scikit-learn."""
 
+    _valid_framework_versions = ["0.20.0"]
+
     def __init__(
         self,
         framework_version,
@@ -84,6 +86,13 @@ def __init__(
         session = sagemaker_session or Session()
         region = session.boto_region_name
 
+        if framework_version not in self._valid_framework_versions:
+            raise ValueError(
+                "scikit-learn version {} is not supported. Supported versions are {}".format(
+                    framework_version, self._valid_framework_versions
+                )
+            )
+
         if not command:
             command = ["python3"]