change: use images_uris.retrieve() for scikit-learn classes

Author: laurenyu

src/sagemaker/image_uri_config/scikit-learn.json

@@ -0,0 +1,62 @@
+{
+    "processors": ["cpu"],
+    "scope": ["inference", "training"],
+    "versions": {
+        "0.20.0": {
+            "py_versions": ["py3"],
+            "registries": {
+                "ap-east-1": "651117190479",
+                "ap-northeast-1": "354813040037",
+                "ap-northeast-2": "366743142698",
+                "ap-south-1": "720646828776",
+                "ap-southeast-1": "121021644041",
+                "ap-southeast-2": "783357654285",
+                "ca-central-1": "341280168497",
+                "cn-north-1": "450853457545",
+                "cn-northwest-1": "451049120500",
+                "eu-central-1": "492215442770",
+                "eu-north-1": "662702820516",
+                "eu-west-1": "141502667606",
+                "eu-west-2": "764974769150",
+                "eu-west-3": "659782779980",
+                "me-south-1": "801668240914",
+                "sa-east-1": "737474898029",
+                "us-east-1": "683313688378",
+                "us-east-2": "257758044811",
+                "us-gov-west-1": "414596584902",
+                "us-iso-east-1": "833128469047",
+                "us-west-1": "746614075791",
+                "us-west-2": "246618743249"
+            },
+            "repository": "sagemaker-scikit-learn"
+        },
+        "0.23-1": {
+            "py_versions": ["py3"],
+            "registries": {
+                "ap-east-1": "651117190479",
+                "ap-northeast-1": "354813040037",
+                "ap-northeast-2": "366743142698",
+                "ap-south-1": "720646828776",
+                "ap-southeast-1": "121021644041",
+                "ap-southeast-2": "783357654285",
+                "ca-central-1": "341280168497",
+                "cn-north-1": "450853457545",
+                "cn-northwest-1": "451049120500",
+                "eu-central-1": "492215442770",
+                "eu-north-1": "662702820516",
+                "eu-west-1": "141502667606",
+                "eu-west-2": "764974769150",
+                "eu-west-3": "659782779980",
+                "me-south-1": "801668240914",
+                "sa-east-1": "737474898029",
+                "us-east-1": "683313688378",
+                "us-east-2": "257758044811",
+                "us-gov-west-1": "414596584902",
+                "us-iso-east-1": "833128469047",
+                "us-west-1": "746614075791",
+                "us-west-2": "246618743249"
+            },
+            "repository": "sagemaker-scikit-learn"
+        }
+    }
+}

src/sagemaker/sklearn/estimator.py

@@ -15,8 +15,8 @@
 
 import logging
 
+from sagemaker import image_uris
 from sagemaker.estimator import Framework
-from sagemaker.fw_registry import default_framework_uri
 from sagemaker.fw_utils import (
     framework_name_from_image,
     framework_version_from_tag,
@@ -137,9 +137,12 @@ def __init__(
         )
 
         if image_uri is None:
-            image_tag = "{}-{}-{}".format(framework_version, "cpu", py_version)
-            self.image_uri = default_framework_uri(
-                SKLearn.__framework_name__, self.sagemaker_session.boto_region_name, image_tag
+            self.image_uri = image_uris.retrieve(
+                SKLearn.__framework_name__,
+                self.sagemaker_session.boto_region_name,
+                version=self.framework_version,
+                py_version=self.py_version,
+                instance_type=instance_type,
             )
 
     def create_model(

src/sagemaker/sklearn/model.py

@@ -16,8 +16,8 @@
 import logging
 
 import sagemaker
+from sagemaker import image_uris
 from sagemaker.deserializers import NumpyDeserializer
-from sagemaker.fw_registry import default_framework_uri
 from sagemaker.fw_utils import model_code_key_prefix, validate_version_or_image_args
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.predictor import Predictor
@@ -163,17 +163,21 @@ def prepare_container_def(self, instance_type=None, accelerator_type=None):
         )
         return sagemaker.container_def(deploy_image, model_data_uri, deploy_env)
 
-    def serving_image_uri(self, region_name, instance_type):  # pylint: disable=unused-argument
+    def serving_image_uri(self, region_name, instance_type):
         """Create a URI for the serving image.
 
         Args:
             region_name (str): AWS region where the image is uploaded.
-            instance_type (str): SageMaker instance type. This parameter is unused because
-                Scikit-learn supports only CPU.
+            instance_type (str): SageMaker instance type.
 
         Returns:
             str: The appropriate image URI based on the given parameters.
 
         """
-        image_tag = "{}-{}-{}".format(self.framework_version, "cpu", self.py_version)
-        return default_framework_uri(self.__framework_name__, region_name, image_tag)
+        return image_uris.retrieve(
+            self.__framework_name__,
+            region_name,
+            version=self.framework_version,
+            py_version=self.py_version,
+            instance_type=instance_type,
+        )

src/sagemaker/sklearn/processing.py

@@ -17,17 +17,14 @@
 """
 from __future__ import absolute_import
 
-from sagemaker.fw_registry import default_framework_uri
-
-from sagemaker import Session
+from sagemaker import image_uris, Session
 from sagemaker.processing import ScriptProcessor
+from sagemaker.sklearn import defaults
 
 
 class SKLearnProcessor(ScriptProcessor):
     """Handles Amazon SageMaker processing tasks for jobs using scikit-learn."""
 
-    _valid_framework_versions = ["0.20.0"]
-
     def __init__(
         self,
         framework_version,
@@ -83,21 +80,14 @@ def __init__(
                 object that configures network isolation, encryption of
                 inter-container traffic, security group IDs, and subnets.
         """
-        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"]
 
-        image_tag = "{}-{}-{}".format(framework_version, "cpu", "py3")
-        image_uri = default_framework_uri("scikit-learn", region, image_tag)
+        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,

tests/conftest.py

@@ -141,11 +141,6 @@ def pytorch_eia_py_version():
     return "py3"
 
 
-@pytest.fixture(scope="module", params=["0.20.0"])
-def sklearn_version(request):
-    return request.param
-
-
 @pytest.fixture(scope="module")
 def xgboost_framework_version(xgboost_version):
     if xgboost_version in ("1", "latest"):
@@ -202,16 +197,6 @@ def rl_ray_full_version():
     return RLEstimator.RAY_LATEST_VERSION
 
 
-@pytest.fixture(scope="module")
-def sklearn_full_version():
-    return "0.20.0"
-
-
-@pytest.fixture(scope="module")
-def sklearn_full_py_version():
-    return "py3"
-
-
 @pytest.fixture(scope="module")
 def tf_full_version(tensorflow_training_latest_version, tensorflow_inference_latest_version):
     """Fixture for TF tests that test both training and inference.
@@ -300,8 +285,9 @@ def pytest_generate_tests(metafunc):
 
 
 def _generate_all_framework_version_fixtures(metafunc):
-    for fw in ("chainer", "mxnet", "pytorch", "tensorflow", "xgboost"):
+    for fw in ("chainer", "mxnet", "pytorch", "scikit-learn", "tensorflow", "xgboost"):
         config = image_uris.config_for_framework(fw)
+        fw = fw.replace("-", "_")  # for fixture names
         if "scope" in config:
             _parametrize_framework_version_fixtures(metafunc, fw, config)
         else:

tests/integ/test_airflow_config.py

@@ -478,7 +478,10 @@ def test_mxnet_airflow_config_uploads_data_source_to_s3(
 
 @pytest.mark.canary_quick
 def test_sklearn_airflow_config_uploads_data_source_to_s3(
-    sagemaker_session, cpu_instance_type, sklearn_full_version, sklearn_full_py_version
+    sagemaker_session,
+    cpu_instance_type,
+    scikit_learn_latest_version,
+    scikit_learn_latest_py_version,
 ):
     with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
         script_path = os.path.join(DATA_DIR, "sklearn_mnist", "mnist.py")
@@ -488,8 +491,8 @@ def test_sklearn_airflow_config_uploads_data_source_to_s3(
             entry_point=script_path,
             role=ROLE,
             instance_type=cpu_instance_type,
-            framework_version=sklearn_full_version,
-            py_version=sklearn_full_py_version,
+            framework_version=scikit_learn_latest_version,
+            py_version=scikit_learn_latest_py_version,
             sagemaker_session=sagemaker_session,
             hyperparameters={"epochs": 1},
         )

tests/integ/test_git.py

@@ -138,7 +138,7 @@ def test_private_github(
 @pytest.mark.local_mode
 @pytest.mark.skip("needs a secure authentication approach")
 def test_private_github_with_2fa(
-    sagemaker_local_session, sklearn_full_version, sklearn_full_py_version
+    sagemaker_local_session, scikit_learn_latest_version, scikit_learn_latest_py_version
 ):
     script_path = "mnist.py"
     data_path = os.path.join(DATA_DIR, "sklearn_mnist")
@@ -155,11 +155,11 @@ def test_private_github_with_2fa(
         entry_point=script_path,
         role="SageMakerRole",
         source_dir=source_dir,
-        py_version=sklearn_full_py_version,
+        py_version=scikit_learn_latest_py_version,
         instance_count=1,
         instance_type="local",
         sagemaker_session=sagemaker_local_session,
-        framework_version=sklearn_full_version,
+        framework_version=scikit_learn_latest_version,
         hyperparameters={"epochs": 1},
         git_config=git_config,
     )
@@ -178,7 +178,7 @@ def test_private_github_with_2fa(
                 model_data,
                 "SageMakerRole",
                 entry_point=script_path,
-                framework_version=sklearn_full_version,
+                framework_version=scikit_learn_latest_version,
                 source_dir=source_dir,
                 sagemaker_session=sagemaker_local_session,
                 git_config=git_config,
@@ -194,7 +194,7 @@ def test_private_github_with_2fa(
 
 @pytest.mark.local_mode
 def test_github_with_ssh_passphrase_not_configured(
-    sagemaker_local_session, sklearn_full_version, sklearn_full_py_version
+    sagemaker_local_session, scikit_learn_latest_version, scikit_learn_latest_py_version
 ):
     script_path = "mnist.py"
     data_path = os.path.join(DATA_DIR, "sklearn_mnist")
@@ -212,8 +212,8 @@ def test_github_with_ssh_passphrase_not_configured(
         instance_count=1,
         instance_type="local",
         sagemaker_session=sagemaker_local_session,
-        framework_version=sklearn_full_version,
-        py_version=sklearn_full_py_version,
+        framework_version=scikit_learn_latest_version,
+        py_version=scikit_learn_latest_py_version,
         hyperparameters={"epochs": 1},
         git_config=git_config,
     )

tests/integ/test_processing.py

@@ -90,12 +90,12 @@ def output_kms_key(sagemaker_session):
     )
 
 
-def test_sklearn(sagemaker_session, sklearn_full_version, cpu_instance_type):
+def test_sklearn(sagemaker_session, scikit_learn_latest_version, cpu_instance_type):
     script_path = os.path.join(DATA_DIR, "dummy_script.py")
     input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 
     sklearn_processor = SKLearnProcessor(
-        framework_version=sklearn_full_version,
+        framework_version=scikit_learn_latest_version,
         role=ROLE,
         instance_type=cpu_instance_type,
         instance_count=1,
@@ -129,12 +129,12 @@ def test_sklearn(sagemaker_session, sklearn_full_version, cpu_instance_type):
 
 @pytest.mark.canary_quick
 def test_sklearn_with_customizations(
-    sagemaker_session, image_uri, sklearn_full_version, cpu_instance_type, output_kms_key
+    sagemaker_session, image_uri, scikit_learn_latest_version, cpu_instance_type, output_kms_key
 ):
     input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 
     sklearn_processor = SKLearnProcessor(
-        framework_version=sklearn_full_version,
+        framework_version=scikit_learn_latest_version,
         role=ROLE,
         command=["python3"],
         instance_type=cpu_instance_type,
@@ -211,14 +211,14 @@ def test_sklearn_with_custom_default_bucket(
     sagemaker_session_with_custom_bucket,
     custom_bucket_name,
     image_uri,
-    sklearn_full_version,
+    scikit_learn_latest_version,
     cpu_instance_type,
     output_kms_key,
 ):
     input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 
     sklearn_processor = SKLearnProcessor(
-        framework_version=sklearn_full_version,
+        framework_version=scikit_learn_latest_version,
         role=ROLE,
         command=["python3"],
         instance_type=cpu_instance_type,
@@ -294,10 +294,10 @@ def test_sklearn_with_custom_default_bucket(
 
 
 def test_sklearn_with_no_inputs_or_outputs(
-    sagemaker_session, image_uri, sklearn_full_version, cpu_instance_type
+    sagemaker_session, image_uri, scikit_learn_latest_version, cpu_instance_type
 ):
     sklearn_processor = SKLearnProcessor(
-        framework_version=sklearn_full_version,
+        framework_version=scikit_learn_latest_version,
         role=ROLE,
         command=["python3"],
         instance_type=cpu_instance_type,
@@ -646,12 +646,14 @@ def test_processor_with_custom_bucket(
     assert job_description["StoppingCondition"] == {"MaxRuntimeInSeconds": 3600}
 
 
-def test_sklearn_with_network_config(sagemaker_session, sklearn_full_version, cpu_instance_type):
+def test_sklearn_with_network_config(
+    sagemaker_session, scikit_learn_latest_version, cpu_instance_type
+):
     script_path = os.path.join(DATA_DIR, "dummy_script.py")
     input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
 
     sklearn_processor = SKLearnProcessor(
-        framework_version=sklearn_full_version,
+        framework_version=scikit_learn_latest_version,
         role=ROLE,
         instance_type=cpu_instance_type,
         instance_count=1,