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

Author: laurenyu

src/sagemaker/image_uri_config/sklearn.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/defaults.py

@@ -13,4 +13,4 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-SKLEARN_NAME = "scikit-learn"
+SKLEARN_NAME = "sklearn"

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(
@@ -243,7 +246,7 @@ class constructor
             init_params["image_uri"] = image_uri
             return init_params
 
-        if framework and framework != cls.__framework_name__:
+        if framework and framework != "scikit-learn":
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
                     job_details["TrainingJobName"]

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,7 +285,7 @@ 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", "sklearn", "tensorflow", "xgboost"):
         config = image_uris.config_for_framework(fw)
         if "scope" in config:
             _parametrize_framework_version_fixtures(metafunc, fw, config)

tests/data/sklearn_mnist/mnist.py

@@ -16,8 +16,8 @@
 import numpy as np
 import os
 
+import joblib
 from sklearn import svm
-from sklearn.externals import joblib
 
 
 def preprocess_mnist(raw, withlabel, ndim, scale, image_dtype, label_dtype, rgb_format):

tests/integ/test_airflow_config.py

@@ -478,7 +478,7 @@ 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, sklearn_latest_version, sklearn_latest_py_version,
 ):
     with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
         script_path = os.path.join(DATA_DIR, "sklearn_mnist", "mnist.py")
@@ -488,8 +488,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=sklearn_latest_version,
+            py_version=sklearn_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, sklearn_latest_version, sklearn_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=sklearn_latest_py_version,
         instance_count=1,
         instance_type="local",
         sagemaker_session=sagemaker_local_session,
-        framework_version=sklearn_full_version,
+        framework_version=sklearn_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=sklearn_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, sklearn_latest_version, sklearn_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=sklearn_latest_version,
+        py_version=sklearn_latest_py_version,
         hyperparameters={"epochs": 1},
         git_config=git_config,
     )