fix: remove py_version from SKLearnProcessor (#285)

Author: knakad

src/sagemaker/sklearn/processing.py

@@ -23,10 +23,6 @@
 from sagemaker.processing import ScriptProcessor
 
 
-_PYTHON_VERSION_TO_COMMAND_MAPPING = {"py2": ["python2"], "py3": ["python3"]}
-_VALID_PYTHON_VERSIONS = ["py2", "py3"]
-
-
 class SKLearnProcessor(ScriptProcessor):
     """Handles Amazon SageMaker processing tasks for jobs using scikit-learn."""
 
@@ -37,7 +33,6 @@ def __init__(
         instance_type,
         instance_count,
         command=None,
-        py_version="py3",
         volume_size_in_gb=30,
         volume_kms_key=None,
         output_kms_key=None,
@@ -65,7 +60,6 @@ def __init__(
             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.
-            py_version (str): The python version to use, for example, 'py3'.
             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
@@ -90,15 +84,10 @@ def __init__(
         session = sagemaker_session or Session()
         region = session.boto_region_name
 
-        if py_version not in _VALID_PYTHON_VERSIONS:
-            raise ValueError(
-                "'py_version' must be a valid value. Please provide one of: 'py2', 'py3'"
-            )
-
         if not command:
-            command = _PYTHON_VERSION_TO_COMMAND_MAPPING[py_version]
+            command = ["python3"]
 
-        image_tag = "{}-{}-{}".format(framework_version, "cpu", py_version)
+        image_tag = "{}-{}-{}".format(framework_version, "cpu", "py3")
         image_uri = default_framework_uri("scikit-learn", region, image_tag)
 
         super(SKLearnProcessor, self).__init__(

tests/integ/test_processing.py

@@ -105,7 +105,6 @@ def test_sklearn_with_customizations(
         command=["python3"],
         instance_type=cpu_instance_type,
         instance_count=1,
-        py_version="py3",
         volume_size_in_gb=100,
         volume_kms_key=None,
         output_kms_key=output_kms_key,
@@ -181,7 +180,6 @@ def test_sklearn_with_no_inputs_or_outputs(
         command=["python3"],
         instance_type=cpu_instance_type,
         instance_count=1,
-        py_version="py3",
         volume_size_in_gb=100,
         volume_kms_key=None,
         max_runtime_in_seconds=3600,

tests/unit/test_processing.py

@@ -171,7 +171,6 @@ def test_sklearn_with_all_customizations(sagemaker_session):
         command=["python3"],
         instance_type="ml.m4.xlarge",
         instance_count=1,
-        py_version="py3",
         volume_size_in_gb=100,
         volume_kms_key=None,
         output_kms_key="arn:aws:kms:us-west-2:012345678901:key/kms-key",
@@ -462,15 +461,3 @@ def test_byo_container_with_baked_in_script(sagemaker_session):
         "experiment_config": None,
     }
     sagemaker_session.process.assert_called_with(**expected_args)
-
-
-def test_sklearn_processor_raises_value_error_if_invalid_py_version_passed_in(sagemaker_session):
-    with pytest.raises(ValueError):
-        SKLearnProcessor(
-            framework_version="0.20.0",
-            role=ROLE,
-            py_version="INVALID_PYTHON_VERSION",
-            instance_type="ml.m4.xlarge",
-            instance_count=1,
-            sagemaker_session=sagemaker_session,
-        )