diff --git a/src/sagemaker/image_uri_config/sklearn.json b/src/sagemaker/image_uri_config/sklearn.json new file mode 100644 index 0000000000..962c2632e6 --- /dev/null +++ b/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" + } + } +} diff --git a/src/sagemaker/sklearn/defaults.py b/src/sagemaker/sklearn/defaults.py index 1d58981fb6..d9f41b029f 100644 --- a/src/sagemaker/sklearn/defaults.py +++ b/src/sagemaker/sklearn/defaults.py @@ -13,4 +13,4 @@ """Placeholder docstring""" from __future__ import absolute_import -SKLEARN_NAME = "scikit-learn" +SKLEARN_NAME = "sklearn" diff --git a/src/sagemaker/sklearn/estimator.py b/src/sagemaker/sklearn/estimator.py index 104682d850..15980b325a 100644 --- a/src/sagemaker/sklearn/estimator.py +++ b/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"] diff --git a/src/sagemaker/sklearn/model.py b/src/sagemaker/sklearn/model.py index e7f5d3e891..7597e7bd53 100644 --- a/src/sagemaker/sklearn/model.py +++ b/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, + ) diff --git a/src/sagemaker/sklearn/processing.py b/src/sagemaker/sklearn/processing.py index 3828cae98b..2722eb2264 100644 --- a/src/sagemaker/sklearn/processing.py +++ b/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, diff --git a/tests/conftest.py b/tests/conftest.py index 0bb36dca06..3d8e586e9c 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/data/sklearn_mnist/mnist.py b/tests/data/sklearn_mnist/mnist.py index 1faf0aa14d..7b8014269b 100644 --- a/tests/data/sklearn_mnist/mnist.py +++ b/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): diff --git a/tests/integ/test_airflow_config.py b/tests/integ/test_airflow_config.py index 8055947f7d..155c03ea66 100644 --- a/tests/integ/test_airflow_config.py +++ b/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}, ) diff --git a/tests/integ/test_git.py b/tests/integ/test_git.py index bb08f74c4a..e2bcdf13b3 100644 --- a/tests/integ/test_git.py +++ b/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, ) diff --git a/tests/integ/test_processing.py b/tests/integ/test_processing.py index 1456f2a0f7..eda8e3445d 100644 --- a/tests/integ/test_processing.py +++ b/tests/integ/test_processing.py @@ -16,9 +16,9 @@ import pytest from botocore.config import Config -from sagemaker import Session -from sagemaker.fw_registry import default_framework_uri +from sagemaker import image_uris, Session +from sagemaker.network import NetworkConfig from sagemaker.processing import ( ProcessingInput, ProcessingOutput, @@ -27,7 +27,6 @@ ProcessingJob, ) from sagemaker.sklearn.processing import SKLearnProcessor -from sagemaker.network import NetworkConfig from tests.integ import DATA_DIR from tests.integ.kms_utils import get_or_create_kms_key @@ -59,10 +58,15 @@ def sagemaker_session_with_custom_bucket( @pytest.fixture(scope="module") -def image_uri(sagemaker_session): - image_tag = "{}-{}-{}".format("0.20.0", "cpu", "py3") - return default_framework_uri( - "scikit-learn", sagemaker_session.boto_session.region_name, image_tag +def image_uri( + sklearn_latest_version, sklearn_latest_py_version, cpu_instance_type, sagemaker_session, +): + return image_uris.retrieve( + "sklearn", + sagemaker_session.boto_region_name, + version=sklearn_latest_version, + py_version=sklearn_latest_py_version, + instance_type=cpu_instance_type, ) @@ -90,12 +94,12 @@ def output_kms_key(sagemaker_session): ) -def test_sklearn(sagemaker_session, sklearn_full_version, cpu_instance_type): +def test_sklearn(sagemaker_session, sklearn_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=sklearn_latest_version, role=ROLE, instance_type=cpu_instance_type, instance_count=1, @@ -129,12 +133,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, sklearn_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=sklearn_latest_version, role=ROLE, command=["python3"], instance_type=cpu_instance_type, @@ -211,14 +215,14 @@ def test_sklearn_with_custom_default_bucket( sagemaker_session_with_custom_bucket, custom_bucket_name, image_uri, - sklearn_full_version, + sklearn_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=sklearn_latest_version, role=ROLE, command=["python3"], instance_type=cpu_instance_type, @@ -294,10 +298,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, sklearn_latest_version, cpu_instance_type ): sklearn_processor = SKLearnProcessor( - framework_version=sklearn_full_version, + framework_version=sklearn_latest_version, role=ROLE, command=["python3"], instance_type=cpu_instance_type, @@ -646,12 +650,12 @@ 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, sklearn_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=sklearn_latest_version, role=ROLE, instance_type=cpu_instance_type, instance_count=1, diff --git a/tests/integ/test_sklearn.py b/tests/integ/test_sklearn.py index d34b5fac86..a98201fd19 100644 --- a/tests/integ/test_sklearn.py +++ b/tests/integ/test_sklearn.py @@ -31,16 +31,16 @@ "This test should be fixed. Details in https://github.com/aws/sagemaker-python-sdk/pull/968" ) def sklearn_training_job( - sagemaker_session, sklearn_full_version, sklearn_full_py_version, cpu_instance_type + sagemaker_session, sklearn_latest_version, sklearn_latest_py_version, cpu_instance_type, ): return _run_mnist_training_job( - sagemaker_session, cpu_instance_type, sklearn_full_version, sklearn_full_py_version + sagemaker_session, cpu_instance_type, sklearn_latest_version, sklearn_latest_py_version, ) sagemaker_session.boto_region_name def test_training_with_additional_hyperparameters( - sagemaker_session, sklearn_full_version, sklearn_full_py_version, cpu_instance_type + sagemaker_session, sklearn_latest_version, sklearn_latest_py_version, cpu_instance_type, ): with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES): script_path = os.path.join(DATA_DIR, "sklearn_mnist", "mnist.py") @@ -50,8 +50,8 @@ def test_training_with_additional_hyperparameters( entry_point=script_path, role="SageMakerRole", 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}, ) @@ -65,11 +65,10 @@ def test_training_with_additional_hyperparameters( job_name = unique_name_from_base("test-sklearn-hp") sklearn.fit({"train": train_input, "test": test_input}, job_name=job_name) - return sklearn.latest_training_job.name def test_training_with_network_isolation( - sagemaker_session, sklearn_full_version, sklearn_full_py_version, cpu_instance_type + sagemaker_session, sklearn_latest_version, sklearn_latest_py_version, cpu_instance_type, ): with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES): script_path = os.path.join(DATA_DIR, "sklearn_mnist", "mnist.py") @@ -79,8 +78,8 @@ def test_training_with_network_isolation( entry_point=script_path, role="SageMakerRole", 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}, enable_network_isolation=True, @@ -98,7 +97,6 @@ def test_training_with_network_isolation( assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[ "EnableNetworkIsolation" ] - return sklearn.latest_training_job.name @pytest.mark.canary_quick @@ -123,8 +121,8 @@ def test_deploy_model( sklearn_training_job, sagemaker_session, cpu_instance_type, - sklearn_full_version, - sklearn_full_py_version, + sklearn_latest_version, + sklearn_latest_py_version, ): endpoint_name = "test-sklearn-deploy-model-{}".format(sagemaker_timestamp()) with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session): @@ -137,7 +135,7 @@ def test_deploy_model( model_data, "SageMakerRole", entry_point=script_path, - framework_version=sklearn_full_version, + framework_version=sklearn_latest_version, sagemaker_session=sagemaker_session, ) predictor = model.deploy(1, cpu_instance_type, endpoint_name=endpoint_name) @@ -149,7 +147,7 @@ def test_deploy_model( "This test should be fixed. Details in https://github.com/aws/sagemaker-python-sdk/pull/968" ) def test_async_fit( - sagemaker_session, cpu_instance_type, sklearn_full_version, sklearn_full_py_version + sagemaker_session, cpu_instance_type, sklearn_latest_version, sklearn_latest_py_version, ): endpoint_name = "test-sklearn-attach-deploy-{}".format(sagemaker_timestamp()) @@ -157,7 +155,7 @@ def test_async_fit( training_job_name = _run_mnist_training_job( sagemaker_session, cpu_instance_type, - sklearn_full_version=sklearn_full_version, + sklearn_version=sklearn_latest_version, wait=False, ) @@ -174,7 +172,7 @@ def test_async_fit( def test_failed_training_job( - sagemaker_session, sklearn_full_version, sklearn_full_py_version, cpu_instance_type + sagemaker_session, sklearn_latest_version, sklearn_latest_py_version, cpu_instance_type, ): with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES): script_path = os.path.join(DATA_DIR, "sklearn_mnist", "failure_script.py") @@ -183,8 +181,8 @@ def test_failed_training_job( sklearn = SKLearn( entry_point=script_path, role="SageMakerRole", - framework_version=sklearn_full_version, - py_version=sklearn_full_py_version, + framework_version=sklearn_latest_version, + py_version=sklearn_latest_py_version, instance_count=1, instance_type=cpu_instance_type, sagemaker_session=sagemaker_session, diff --git a/tests/unit/sagemaker/image_uris/test_sklearn.py b/tests/unit/sagemaker/image_uris/test_sklearn.py new file mode 100644 index 0000000000..75b213d158 --- /dev/null +++ b/tests/unit/sagemaker/image_uris/test_sklearn.py @@ -0,0 +1,85 @@ +# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +from __future__ import absolute_import + +import pytest + +from sagemaker import image_uris +from tests.unit.sagemaker.image_uris import expected_uris, regions + +ACCOUNTS = { + "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", +} + + +def test_valid_uris(sklearn_version): + for region in regions.regions(): + uri = image_uris.retrieve( + "sklearn", + region=region, + version=sklearn_version, + py_version="py3", + instance_type="ml.c4.xlarge", + ) + + expected = expected_uris.framework_uri( + "sagemaker-scikit-learn", + sklearn_version, + ACCOUNTS[region], + py_version="py3", + region=region, + ) + assert expected == uri + + +def test_py2_error(sklearn_version): + with pytest.raises(ValueError) as e: + image_uris.retrieve( + "sklearn", + region="us-west-2", + version=sklearn_version, + py_version="py2", + instance_type="ml.c4.xlarge", + ) + + assert "Unsupported Python version: py2." in str(e.value) + + +def test_gpu_error(sklearn_version): + with pytest.raises(ValueError) as e: + image_uris.retrieve( + "sklearn", region="us-west-2", version=sklearn_version, instance_type="ml.p2.xlarge", + ) + + assert "Unsupported processor: gpu." in str(e.value) diff --git a/tests/unit/test_fw_registry.py b/tests/unit/test_fw_registry.py index 8a27fc513e..c5efdcde41 100644 --- a/tests/unit/test_fw_registry.py +++ b/tests/unit/test_fw_registry.py @@ -15,10 +15,9 @@ import pytest from sagemaker.fw_registry import registry, default_framework_uri -from sagemaker.sklearn import SKLearn -scikit_learn_framework_name = SKLearn.__framework_name__ +scikit_learn_framework_name = "scikit-learn" def test_registry_sparkml_serving(): diff --git a/tests/unit/test_processing.py b/tests/unit/test_processing.py index 860b664df9..b0c50ed937 100644 --- a/tests/unit/test_processing.py +++ b/tests/unit/test_processing.py @@ -28,7 +28,7 @@ BUCKET_NAME = "mybucket" REGION = "us-west-2" ROLE = "arn:aws:iam::012345678901:role/SageMakerRole" -ECR_PREFIX = "246618743249.dkr.ecr.us-west-2.amazonaws.com" +ECR_HOSTNAME = "ecr.us-west-2.amazonaws.com" CUSTOM_IMAGE_URI = "012345678901.dkr.ecr.us-west-2.amazonaws.com/my-custom-image-uri" PROCESSING_JOB_DESCRIPTION = { @@ -95,16 +95,18 @@ def sagemaker_session(): return session_mock -@patch("sagemaker.fw_registry.get_ecr_image_uri_prefix", return_value=ECR_PREFIX) +@patch("sagemaker.utils._botocore_resolver") @patch("os.path.exists", return_value=True) @patch("os.path.isfile", return_value=True) def test_sklearn_processor_with_required_parameters( - exists_mock, isfile_mock, ecr_prefix, sagemaker_session + exists_mock, isfile_mock, botocore_resolver, sagemaker_session, sklearn_version ): + botocore_resolver.return_value.construct_endpoint.return_value = {"hostname": ECR_HOSTNAME} + processor = SKLearnProcessor( role=ROLE, instance_type="ml.m4.xlarge", - framework_version="0.20.0", + framework_version=sklearn_version, instance_count=1, sagemaker_session=sagemaker_session, ) @@ -114,20 +116,24 @@ def test_sklearn_processor_with_required_parameters( expected_args = _get_expected_args(processor._current_job_name) sklearn_image_uri = ( - "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3" - ) + "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:{}-cpu-py3" + ).format(sklearn_version) expected_args["app_specification"]["ImageUri"] = sklearn_image_uri sagemaker_session.process.assert_called_with(**expected_args) -@patch("sagemaker.fw_registry.get_ecr_image_uri_prefix", return_value=ECR_PREFIX) +@patch("sagemaker.utils._botocore_resolver") @patch("os.path.exists", return_value=True) @patch("os.path.isfile", return_value=True) -def test_sklearn_with_all_parameters(exists_mock, isfile_mock, ecr_prefix, sagemaker_session): +def test_sklearn_with_all_parameters( + exists_mock, isfile_mock, botocore_resolver, sklearn_version, sagemaker_session +): + botocore_resolver.return_value.construct_endpoint.return_value = {"hostname": ECR_HOSTNAME} + processor = SKLearnProcessor( role=ROLE, - framework_version="0.20.0", + framework_version=sklearn_version, instance_type="ml.m4.xlarge", instance_count=1, volume_size_in_gb=100, @@ -176,28 +182,13 @@ def test_sklearn_with_all_parameters(exists_mock, isfile_mock, ecr_prefix, sagem expected_args = _get_expected_args_all_parameters(processor._current_job_name) sklearn_image_uri = ( - "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3" - ) + "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:{}-cpu-py3" + ).format(sklearn_version) expected_args["app_specification"]["ImageUri"] = sklearn_image_uri sagemaker_session.process.assert_called_with(**expected_args) -@patch("os.path.exists", return_value=True) -@patch("os.path.isfile", return_value=True) -def test_sklearn_processor_errors_with_invalid_framework_version( - exists_mock, isfile_mock, sagemaker_session -): - with pytest.raises(ValueError): - SKLearnProcessor( - role=ROLE, - framework_version="0.21.0", - instance_type="ml.m4.xlarge", - instance_count=1, - sagemaker_session=sagemaker_session, - ) - - @patch("os.path.exists", return_value=False) def test_script_processor_errors_with_nonexistent_local_code(exists_mock, sagemaker_session): processor = _get_script_processor(sagemaker_session) diff --git a/tests/unit/test_sklearn.py b/tests/unit/test_sklearn.py index 3723616bca..57925c70d7 100644 --- a/tests/unit/test_sklearn.py +++ b/tests/unit/test_sklearn.py @@ -156,11 +156,7 @@ def test_train_image(sagemaker_session, sklearn_version): source_dir=source_dir, ) - train_image = sklearn.train_image() - assert ( - train_image - == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:0.20.0-cpu-py3" - ) + assert _get_full_cpu_image_uri(sklearn_version) == sklearn.train_image() def test_create_model(sagemaker_session, sklearn_version): @@ -391,7 +387,7 @@ def test_fail_distributed_training(sagemaker_session, sklearn_version): assert "Scikit-Learn does not support distributed training." in str(error) -def test_fail_GPU_training(sagemaker_session, sklearn_version): +def test_fail_gpu_training(sagemaker_session, sklearn_version): with pytest.raises(ValueError) as error: SKLearn( entry_point=SCRIPT_PATH, @@ -416,30 +412,6 @@ def test_model(sagemaker_session, sklearn_version): assert isinstance(predictor, SKLearnPredictor) -def test_train_image_default(sagemaker_session, sklearn_version): - sklearn = SKLearn( - entry_point=SCRIPT_PATH, - framework_version=sklearn_version, - py_version=PYTHON_VERSION, - role=ROLE, - sagemaker_session=sagemaker_session, - instance_type=INSTANCE_TYPE, - ) - - assert _get_full_cpu_image_uri(sklearn_version) in sklearn.train_image() - - -def test_train_image_cpu_instances(sagemaker_session, sklearn_version): - sklearn = _sklearn_estimator(sagemaker_session, sklearn_version, instance_type="ml.c2.2xlarge") - assert sklearn.train_image() == _get_full_cpu_image_uri(sklearn_version) - - sklearn = _sklearn_estimator(sagemaker_session, sklearn_version, instance_type="ml.c4.2xlarge") - assert sklearn.train_image() == _get_full_cpu_image_uri(sklearn_version) - - sklearn = _sklearn_estimator(sagemaker_session, sklearn_version, instance_type="ml.m16") - assert sklearn.train_image() == _get_full_cpu_image_uri(sklearn_version) - - def test_attach(sagemaker_session, sklearn_version): training_image = "1.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:{}-cpu-{}".format( sklearn_version, PYTHON_VERSION @@ -588,11 +560,3 @@ def test_model_py2_raises(sagemaker_session, sklearn_version): framework_version=sklearn_version, py_version="py2", ) - - -def test_custom_image_estimator_deploy(sagemaker_session, sklearn_version): - custom_image = "mycustomimage:latest" - sklearn = _sklearn_estimator(sagemaker_session, sklearn_version) - sklearn.fit(inputs="s3://mybucket/train", job_name="new_name") - model = sklearn.create_model(image_uri=custom_image) - assert model.image_uri == custom_image