Fix SKLearnModel default account in image uri. (#624)

Author: ericangelokim

src/sagemaker/sklearn/defaults.py

@@ -12,4 +12,6 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+SKLEARN_NAME = 'scikit-learn'
+
 SKLEARN_VERSION = '0.20.0'

src/sagemaker/sklearn/estimator.py

@@ -17,7 +17,7 @@
 from sagemaker.estimator import Framework
 from sagemaker.fw_registry import default_framework_uri
 from sagemaker.fw_utils import framework_name_from_image, empty_framework_version_warning
-from sagemaker.sklearn.defaults import SKLEARN_VERSION
+from sagemaker.sklearn.defaults import SKLEARN_VERSION, SKLEARN_NAME
 from sagemaker.sklearn.model import SKLearnModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -28,7 +28,7 @@
 class SKLearn(Framework):
     """Handle end-to-end training and deployment of custom Scikit-learn code."""
 
-    __framework_name__ = "scikit-learn"
+    __framework_name__ = SKLEARN_NAME
 
     def __init__(self, entry_point, framework_version=SKLEARN_VERSION, source_dir=None, hyperparameters=None,
                  py_version='py3', image_name=None, **kwargs):
@@ -74,7 +74,7 @@ def __init__(self, entry_point, framework_version=SKLEARN_VERSION, source_dir=No
         train_instance_count = kwargs.get('train_instance_count')
         if train_instance_count:
             if train_instance_count != 1:
-                raise AttributeError("SciKit-Learn does not support distributed training. "
+                raise AttributeError("Scikit-Learn does not support distributed training. "
                                      "Please remove the 'train_instance_count' argument or set "
                                      "'train_instance_count=1' when initializing SKLearn.")
         super(SKLearn, self).__init__(entry_point, source_dir, hyperparameters, image_name=image_name,
@@ -154,6 +154,6 @@ def _validate_not_gpu_instance_type(training_instance_type):
                           'ml.p3.xlarge', 'ml.p3.8xlarge', 'ml.p3.16xlarge']
 
     if training_instance_type in gpu_instance_types:
-        raise ValueError("GPU training in not supported for SciKit-Learn. "
+        raise ValueError("GPU training in not supported for Scikit-Learn. "
                          "Please pick a different instance type from here: "
                          "https://aws.amazon.com/ec2/instance-types/")

src/sagemaker/sklearn/model.py

@@ -13,10 +13,11 @@
 from __future__ import absolute_import
 
 import sagemaker
-from sagemaker.fw_utils import create_image_uri, model_code_key_prefix
+from sagemaker.fw_utils import model_code_key_prefix
+from sagemaker.fw_registry import default_framework_uri
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.predictor import RealTimePredictor, npy_serializer, numpy_deserializer
-from sagemaker.sklearn.defaults import SKLEARN_VERSION
+from sagemaker.sklearn.defaults import SKLEARN_VERSION, SKLEARN_NAME
 
 
 class SKLearnPredictor(RealTimePredictor):
@@ -40,7 +41,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
 class SKLearnModel(FrameworkModel):
     """An Scikit-learn SageMaker ``Model`` that can be deployed to a SageMaker ``Endpoint``."""
 
-    __framework_name__ = 'scikit-learn'
+    __framework_name__ = SKLEARN_NAME
 
     def __init__(self, model_data, role, entry_point, image=None, py_version='py3', framework_version=SKLEARN_VERSION,
                  predictor_cls=SKLearnPredictor, model_server_workers=None, **kwargs):
@@ -77,16 +78,22 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
         Args:
             instance_type (str): The EC2 instance type to deploy this Model to. For example, 'ml.p2.xlarge'.
             accelerator_type (str): The Elastic Inference accelerator type to deploy to the instance for loading and
-                making inferences to the model. For example, 'ml.eia1.medium'.
+                making inferences to the model. For example, 'ml.eia1.medium'. Note: accelerator types are not
+                supported by SKLearnModel.
 
         Returns:
             dict[str, str]: A container definition object usable with the CreateModel API.
         """
+        if accelerator_type:
+            raise ValueError("Accelerator types are not supported for Scikit-Learn.")
+
         deploy_image = self.image
         if not deploy_image:
-            region_name = self.sagemaker_session.boto_session.region_name
-            deploy_image = create_image_uri(region_name, self.__framework_name__, instance_type,
-                                            self.framework_version, self.py_version, accelerator_type=accelerator_type)
+            image_tag = "{}-{}-{}".format(self.framework_version, "cpu", self.py_version)
+            deploy_image = default_framework_uri(
+                self.__framework_name__,
+                self.sagemaker_session.boto_region_name,
+                image_tag)
 
         deploy_key_prefix = model_code_key_prefix(self.key_prefix, self.name, deploy_image)
         self._upload_code(deploy_key_prefix)

tests/unit/test_sklearn.py

@@ -128,7 +128,17 @@ def test_train_image(sagemaker_session, sklearn_version):
     assert train_image == '246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3'
 
 
-def test_create_model(sagemaker_session, sklearn_version):
+def test_create_model(sagemaker_session):
+    source_dir = 's3://mybucket/source'
+
+    sklearn_model = SKLearnModel(model_data=source_dir, role=ROLE, sagemaker_session=sagemaker_session,
+                                 entry_point=SCRIPT_PATH)
+    default_image_uri = _get_full_cpu_image_uri('0.20.0')
+    model_values = sklearn_model.prepare_container_def(CPU)
+    assert model_values['Image'] == default_image_uri
+
+
+def test_create_model_from_estimator(sagemaker_session, sklearn_version):
     container_log_level = '"logging.INFO"'
     source_dir = 's3://mybucket/source'
     sklearn = SKLearn(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
@@ -231,15 +241,15 @@ def test_fail_distributed_training(sagemaker_session, sklearn_version):
         SKLearn(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                 train_instance_count=DIST_INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
                 py_version=PYTHON_VERSION, framework_version=sklearn_version)
-    assert "SciKit-Learn does not support distributed training." in str(error)
+    assert "Scikit-Learn does not support distributed training." in str(error)
 
 
 def test_fail_GPU_training(sagemaker_session, sklearn_version):
     with pytest.raises(ValueError) as error:
         SKLearn(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                 train_instance_type=GPU_INSTANCE_TYPE, py_version=PYTHON_VERSION,
                 framework_version=sklearn_version)
-    assert "GPU training in not supported for SciKit-Learn." in str(error)
+    assert "GPU training in not supported for Scikit-Learn." in str(error)
 
 
 def test_model(sagemaker_session):