Re-enable the new SKLearnProcessor

Author: Verdi March

src/sagemaker/sklearn/processing.py

@@ -17,93 +17,67 @@
 """
 from __future__ import absolute_import
 
-from sagemaker import image_uris, Session
-from sagemaker.processing import ScriptProcessor
-from sagemaker.sklearn import defaults
+from sagemaker.processing import FrameworkProcessor
+from sagemaker.sklearn.estimator import SKLearn
 
 
-class SKLearnProcessor(ScriptProcessor):
-    """Handles Amazon SageMaker processing tasks for jobs using scikit-learn."""
+class SKLearnProcessor(FrameworkProcessor):
+    """Initialize an ``SKLearnProcessor`` instance.
+
+    The SKLearnProcessor handles Amazon SageMaker processing tasks for jobs using scikit-learn.
+
+    Unless ``image_uri`` is specified, the scikit-learn environment is an
+    Amazon-built Docker container that executes functions defined in the supplied
+    ``code`` Python script.
+
+    The arguments have the exact same meaning as in ``FrameworkProcessor``.
+
+    .. tip::
+
+        You can find additional parameters for initializing this class at
+        :class:`~sagemaker.processing.FrameworkProcessor`.
+    """
+
+    estimator_cls = SKLearn
 
     def __init__(
         self,
-        framework_version,
+        framework_version,  # New arg
         role,
-        instance_type,
         instance_count,
+        instance_type,
+        py_version="py3",  # New kwarg
+        image_uri=None,
         command=None,
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
+        code_location=None,  # New arg
         max_runtime_in_seconds=None,
         base_job_name=None,
         sagemaker_session=None,
         env=None,
         tags=None,
         network_config=None,
     ):
-        """Initialize an ``SKLearnProcessor`` instance.
-
-        The SKLearnProcessor handles Amazon SageMaker processing tasks for jobs using scikit-learn.
-
-        Args:
-            framework_version (str): The version of scikit-learn.
-            role (str): An AWS IAM role name or ARN. The Amazon SageMaker training jobs
-                and APIs that create Amazon SageMaker endpoints use this role
-                to access training data and model artifacts. After the endpoint
-                is created, the inference code might use the IAM role, if it
-                needs to access an AWS resource.
-            instance_type (str): Type of EC2 instance to use for
-                processing, for example, 'ml.c4.xlarge'.
-            instance_count (int): The number of instances to run
-                the Processing job with. Defaults to 1.
-            command ([str]): The command to run, along with any command-line flags.
-                Example: ["python3", "-v"]. If not provided, ["python3"] or ["python2"]
-                will be chosen based on the py_version parameter.
-            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.
-            output_kms_key (str): The KMS key id for all ProcessingOutputs.
-            max_runtime_in_seconds (int): Timeout in seconds.
-                After this amount of time Amazon SageMaker terminates the job
-                regardless of its current status.
-            base_job_name (str): Prefix for processing name. If not specified,
-                the processor generates a default job name, based on the
-                training image name and current timestamp.
-            sagemaker_session (sagemaker.session.Session): Session object which
-                manages interactions with Amazon SageMaker APIs and any other
-                AWS services needed. If not specified, the processor creates one
-                using the default AWS configuration chain.
-            env (dict): Environment variables to be passed to the processing job.
-            tags ([dict]): List of tags to be passed to the processing job.
-            network_config (sagemaker.network.NetworkConfig): A NetworkConfig
-                object that configures network isolation, encryption of
-                inter-container traffic, security group IDs, and subnets.
-        """
-        if not command:
-            command = ["python3"]
-
-        session = sagemaker_session or Session()
-        region = session.boto_region_name
-
-        image_uri = image_uris.retrieve(
-            defaults.SKLEARN_NAME, region, version=framework_version, instance_type=instance_type
-        )
-
-        super(SKLearnProcessor, self).__init__(
-            role=role,
-            image_uri=image_uri,
-            instance_count=instance_count,
-            instance_type=instance_type,
-            command=command,
-            volume_size_in_gb=volume_size_in_gb,
-            volume_kms_key=volume_kms_key,
-            output_kms_key=output_kms_key,
-            max_runtime_in_seconds=max_runtime_in_seconds,
-            base_job_name=base_job_name,
-            sagemaker_session=session,
-            env=env,
-            tags=tags,
-            network_config=network_config,
+        """This processor executes a Python script in a scikit-learn execution environment."""
+        super().__init__(
+            self.estimator_cls,
+            framework_version,
+            role,
+            instance_count,
+            instance_type,
+            py_version,
+            image_uri,
+            command,
+            volume_size_in_gb,
+            volume_kms_key,
+            output_kms_key,
+            code_location,
+            max_runtime_in_seconds,
+            base_job_name,
+            sagemaker_session,
+            env,
+            tags,
+            network_config,
         )

tests/integ/test_local_mode.py

@@ -349,12 +349,12 @@ def test_local_processing_sklearn(sagemaker_local_session_no_local_code, sklearn
 
     job_description = sklearn_processor.latest_job.describe()
 
-    assert len(job_description["ProcessingInputs"]) == 2
+    assert len(job_description["ProcessingInputs"]) == 3
     assert job_description["ProcessingResources"]["ClusterConfig"]["InstanceCount"] == 1
     assert job_description["ProcessingResources"]["ClusterConfig"]["InstanceType"] == "local"
     assert job_description["AppSpecification"]["ContainerEntrypoint"] == [
         "python3",
-        "/opt/ml/processing/input/code/dummy_script.py",
+        "/opt/ml/processing/input/entrypoint/runproc.py",
     ]
     assert job_description["RoleArn"] == "<no_role>"
 

tests/integ/test_processing.py

@@ -139,7 +139,7 @@ def test_sklearn(sagemaker_session, sklearn_latest_version, cpu_instance_type):
 
     job_description = sklearn_processor.latest_job.describe()
 
-    assert len(job_description["ProcessingInputs"]) == 2
+    assert len(job_description["ProcessingInputs"]) == 3
     assert job_description["ProcessingResources"]["ClusterConfig"]["InstanceCount"] == 1
     assert (
         job_description["ProcessingResources"]["ClusterConfig"]["InstanceType"] == cpu_instance_type
@@ -148,7 +148,7 @@ def test_sklearn(sagemaker_session, sklearn_latest_version, cpu_instance_type):
     assert job_description["StoppingCondition"] == {"MaxRuntimeInSeconds": 86400}
     assert job_description["AppSpecification"]["ContainerEntrypoint"] == [
         "python3",
-        "/opt/ml/processing/input/code/dummy_script.py",
+        "/opt/ml/processing/input/entrypoint/runproc.py",
     ]
     assert ROLE in job_description["RoleArn"]
 
@@ -204,6 +204,7 @@ def test_sklearn_with_customizations(
     assert job_description["ProcessingInputs"][0]["InputName"] == "dummy_input"
 
     assert job_description["ProcessingInputs"][1]["InputName"] == "code"
+    assert job_description["ProcessingInputs"][2]["InputName"] == "entrypoint"
 
     assert job_description["ProcessingJobName"].startswith("test-sklearn-with-customizations")
 
@@ -221,7 +222,7 @@ def test_sklearn_with_customizations(
     assert job_description["AppSpecification"]["ContainerArguments"] == ["-v"]
     assert job_description["AppSpecification"]["ContainerEntrypoint"] == [
         "python3",
-        "/opt/ml/processing/input/code/dummy_script.py",
+        "/opt/ml/processing/input/entrypoint/runproc.py",
     ]
     assert job_description["AppSpecification"]["ImageUri"] == image_uri
 
@@ -288,6 +289,9 @@ def test_sklearn_with_custom_default_bucket(
     assert job_description["ProcessingInputs"][0]["InputName"] == "dummy_input"
     assert custom_bucket_name in job_description["ProcessingInputs"][0]["S3Input"]["S3Uri"]
 
+    assert job_description["ProcessingInputs"][1]["InputName"] == "code"
+    assert custom_bucket_name in job_description["ProcessingInputs"][1]["S3Input"]["S3Uri"]
+
     assert job_description["ProcessingInputs"][2]["InputName"] == "entrypoint"
     assert custom_bucket_name in job_description["ProcessingInputs"][2]["S3Input"]["S3Uri"]
 
@@ -307,7 +311,7 @@ def test_sklearn_with_custom_default_bucket(
     assert job_description["AppSpecification"]["ContainerArguments"] == ["-v"]
     assert job_description["AppSpecification"]["ContainerEntrypoint"] == [
         "python3",
-        "/opt/ml/processing/input/code/dummy_script.py",
+        "/opt/ml/processing/input/entrypoint/runproc.py",
     ]
     assert job_description["AppSpecification"]["ImageUri"] == image_uri
 
@@ -343,6 +347,7 @@ def test_sklearn_with_no_inputs_or_outputs(
     job_description = sklearn_processor.latest_job.describe()
 
     assert job_description["ProcessingInputs"][0]["InputName"] == "code"
+    assert job_description["ProcessingInputs"][1]["InputName"] == "entrypoint"
 
     assert job_description["ProcessingJobName"].startswith("test-sklearn-with-no-inputs")
 
@@ -357,7 +362,7 @@ def test_sklearn_with_no_inputs_or_outputs(
     assert job_description["AppSpecification"]["ContainerArguments"] == ["-v"]
     assert job_description["AppSpecification"]["ContainerEntrypoint"] == [
         "python3",
-        "/opt/ml/processing/input/code/dummy_script.py",
+        "/opt/ml/processing/input/entrypoint/runproc.py",
     ]
     assert job_description["AppSpecification"]["ImageUri"] == image_uri
 

tests/integ/test_sklearn.py

@@ -46,6 +46,20 @@ def sklearn_training_job(
     sagemaker_session.boto_region_name
 
 
+def test_framework_processing_job_with_deps(
+    sagemaker_session,
+    sklearn_latest_version,
+    sklearn_latest_py_version,
+    cpu_instance_type,
+):
+    return _run_processing_job(
+        sagemaker_session,
+        cpu_instance_type,
+        sklearn_latest_version,
+        sklearn_latest_py_version,
+    )
+
+
 def test_training_with_additional_hyperparameters(
     sagemaker_session,
     sklearn_latest_version,