improve imports and error message

Author: Chuyang Deng

src/sagemaker/chainer/estimator.py

@@ -22,7 +22,7 @@
     empty_framework_version_warning,
     python_deprecation_warning,
 )
-from sagemaker.chainer.defaults import CHAINER_VERSION, LATEST_VERSION, LATEST_PY2_VERSION
+from sagemaker.chainer import defaults
 from sagemaker.chainer.model import ChainerModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -40,7 +40,7 @@ class Chainer(Framework):
     _process_slots_per_host = "sagemaker_process_slots_per_host"
     _additional_mpi_options = "sagemaker_additional_mpi_options"
 
-    LATEST_VERSION = LATEST_VERSION
+    LATEST_VERSION = defaults.LATEST_VERSION
 
     def __init__(
         self,
@@ -126,15 +126,19 @@ def __init__(
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
         if framework_version is None:
-            logger.warning(empty_framework_version_warning(CHAINER_VERSION, self.LATEST_VERSION))
-        self.framework_version = framework_version or CHAINER_VERSION
+            logger.warning(
+                empty_framework_version_warning(defaults.CHAINER_VERSION, self.LATEST_VERSION)
+            )
+        self.framework_version = framework_version or defaults.CHAINER_VERSION
 
         super(Chainer, self).__init__(
             entry_point, source_dir, hyperparameters, image_name=image_name, **kwargs
         )
 
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         self.py_version = py_version
         self.use_mpi = use_mpi

src/sagemaker/chainer/model.py

@@ -23,7 +23,7 @@
     empty_framework_version_warning,
 )
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
-from sagemaker.chainer.defaults import CHAINER_VERSION, LATEST_VERSION, LATEST_PY2_VERSION
+from sagemaker.chainer import defaults
 from sagemaker.predictor import RealTimePredictor, npy_serializer, numpy_deserializer
 
 logger = logging.getLogger("sagemaker")
@@ -111,13 +111,17 @@ def __init__(
             model_data, image, role, entry_point, predictor_cls=predictor_cls, **kwargs
         )
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         if framework_version is None:
-            logger.warning(empty_framework_version_warning(CHAINER_VERSION, LATEST_VERSION))
+            logger.warning(
+                empty_framework_version_warning(defaults.CHAINER_VERSION, defaults.LATEST_VERSION)
+            )
 
         self.py_version = py_version
-        self.framework_version = framework_version or CHAINER_VERSION
+        self.framework_version = framework_version or defaults.CHAINER_VERSION
         self.model_server_workers = model_server_workers
 
     def prepare_container_def(self, instance_type, accelerator_type=None):

src/sagemaker/mxnet/estimator.py

@@ -23,7 +23,7 @@
     python_deprecation_warning,
     is_version_equal_or_higher,
 )
-from sagemaker.mxnet.defaults import MXNET_VERSION, LATEST_VERSION, LATEST_PY2_VERSION
+from sagemaker.mxnet import defaults
 from sagemaker.mxnet.model import MXNetModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -36,7 +36,7 @@ class MXNet(Framework):
     __framework_name__ = "mxnet"
     _LOWEST_SCRIPT_MODE_VERSION = ["1", "3"]
 
-    LATEST_VERSION = LATEST_VERSION
+    LATEST_VERSION = defaults.LATEST_VERSION
 
     def __init__(
         self,
@@ -107,8 +107,10 @@ def __init__(
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
         if framework_version is None:
-            logger.warning(empty_framework_version_warning(MXNET_VERSION, self.LATEST_VERSION))
-        self.framework_version = framework_version or MXNET_VERSION
+            logger.warning(
+                empty_framework_version_warning(defaults.MXNET_VERSION, self.LATEST_VERSION)
+            )
+        self.framework_version = framework_version or defaults.MXNET_VERSION
 
         if "enable_sagemaker_metrics" not in kwargs:
             # enable sagemaker metrics for MXNet v1.6 or greater:
@@ -120,7 +122,9 @@ def __init__(
         )
 
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         self.py_version = py_version
         self._configure_distribution(distributions)

src/sagemaker/mxnet/model.py

@@ -25,7 +25,7 @@
     empty_framework_version_warning,
 )
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
-from sagemaker.mxnet.defaults import MXNET_VERSION, LATEST_VERSION, LATEST_PY2_VERSION
+from sagemaker.mxnet import defaults
 from sagemaker.predictor import RealTimePredictor, json_serializer, json_deserializer
 
 logger = logging.getLogger("sagemaker")
@@ -113,13 +113,17 @@ def __init__(
         )
 
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         if framework_version is None:
-            logger.warning(empty_framework_version_warning(MXNET_VERSION, LATEST_VERSION))
+            logger.warning(
+                empty_framework_version_warning(defaults.MXNET_VERSION, defaults.LATEST_VERSION)
+            )
 
         self.py_version = py_version
-        self.framework_version = framework_version or MXNET_VERSION
+        self.framework_version = framework_version or defaults.MXNET_VERSION
         self.model_server_workers = model_server_workers
 
     def prepare_container_def(self, instance_type, accelerator_type=None):

src/sagemaker/pytorch/estimator.py

@@ -23,12 +23,7 @@
     python_deprecation_warning,
     is_version_equal_or_higher,
 )
-from sagemaker.pytorch.defaults import (
-    PYTORCH_VERSION,
-    PYTHON_VERSION,
-    LATEST_VERSION,
-    LATEST_PY2_VERSION,
-)
+from sagemaker.pytorch import defaults
 from sagemaker.pytorch.model import PyTorchModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -40,14 +35,14 @@ class PyTorch(Framework):
 
     __framework_name__ = "pytorch"
 
-    LATEST_VERSION = LATEST_VERSION
+    LATEST_VERSION = defaults.LATEST_VERSION
 
     def __init__(
         self,
         entry_point,
         source_dir=None,
         hyperparameters=None,
-        py_version=PYTHON_VERSION,
+        py_version=defaults.PYTHON_VERSION,
         framework_version=None,
         image_name=None,
         **kwargs
@@ -108,8 +103,10 @@ def __init__(
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
         if framework_version is None:
-            logger.warning(empty_framework_version_warning(PYTORCH_VERSION, self.LATEST_VERSION))
-        self.framework_version = framework_version or PYTORCH_VERSION
+            logger.warning(
+                empty_framework_version_warning(defaults.PYTORCH_VERSION, self.LATEST_VERSION)
+            )
+        self.framework_version = framework_version or defaults.PYTORCH_VERSION
 
         if "enable_sagemaker_metrics" not in kwargs:
             # enable sagemaker metrics for PT v1.3 or greater:
@@ -121,7 +118,9 @@ def __init__(
         )
 
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         self.py_version = py_version
 

src/sagemaker/pytorch/model.py

@@ -24,12 +24,7 @@
     empty_framework_version_warning,
 )
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
-from sagemaker.pytorch.defaults import (
-    PYTORCH_VERSION,
-    PYTHON_VERSION,
-    LATEST_VERSION,
-    LATEST_PY2_VERSION,
-)
+from sagemaker.pytorch import defaults
 from sagemaker.predictor import RealTimePredictor, npy_serializer, numpy_deserializer
 
 logger = logging.getLogger("sagemaker")
@@ -72,7 +67,7 @@ def __init__(
         role,
         entry_point,
         image=None,
-        py_version=PYTHON_VERSION,
+        py_version=defaults.PYTHON_VERSION,
         framework_version=None,
         predictor_cls=PyTorchPredictor,
         model_server_workers=None,
@@ -119,13 +114,17 @@ def __init__(
         )
 
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         if framework_version is None:
-            logger.warning(empty_framework_version_warning(PYTORCH_VERSION, LATEST_VERSION))
+            logger.warning(
+                empty_framework_version_warning(defaults.PYTORCH_VERSION, defaults.LATEST_VERSION)
+            )
 
         self.py_version = py_version
-        self.framework_version = framework_version or PYTORCH_VERSION
+        self.framework_version = framework_version or defaults.PYTORCH_VERSION
         self.model_server_workers = model_server_workers
 
     def prepare_container_def(self, instance_type, accelerator_type=None):

src/sagemaker/sklearn/estimator.py

@@ -22,7 +22,7 @@
     empty_framework_version_warning,
     python_deprecation_warning,
 )
-from sagemaker.sklearn.defaults import SKLEARN_VERSION, SKLEARN_NAME, LATEST_PY2_VERSION
+from sagemaker.sklearn import defaults
 from sagemaker.sklearn.model import SKLearnModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -32,12 +32,12 @@
 class SKLearn(Framework):
     """Handle end-to-end training and deployment of custom Scikit-learn code."""
 
-    __framework_name__ = SKLEARN_NAME
+    __framework_name__ = defaults.SKLEARN_NAME
 
     def __init__(
         self,
         entry_point,
-        framework_version=SKLEARN_VERSION,
+        framework_version=defaults.SKLEARN_VERSION,
         source_dir=None,
         hyperparameters=None,
         py_version="py3",
@@ -119,13 +119,17 @@ def __init__(
         )
 
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         self.py_version = py_version
 
         if framework_version is None:
-            logger.warning(empty_framework_version_warning(SKLEARN_VERSION, SKLEARN_VERSION))
-        self.framework_version = framework_version or SKLEARN_VERSION
+            logger.warning(
+                empty_framework_version_warning(defaults.SKLEARN_VERSION, defaults.SKLEARN_VERSION)
+            )
+        self.framework_version = framework_version or defaults.SKLEARN_VERSION
 
         if image_name is None:
             image_tag = "{}-{}-{}".format(framework_version, "cpu", py_version)

src/sagemaker/sklearn/model.py

@@ -20,7 +20,7 @@
 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, SKLEARN_NAME, LATEST_PY2_VERSION
+from sagemaker.sklearn import defaults
 
 logger = logging.getLogger("sagemaker")
 
@@ -53,7 +53,7 @@ class SKLearnModel(FrameworkModel):
     ``Endpoint``.
     """
 
-    __framework_name__ = SKLEARN_NAME
+    __framework_name__ = defaults.SKLEARN_NAME
 
     def __init__(
         self,
@@ -62,7 +62,7 @@ def __init__(
         entry_point,
         image=None,
         py_version="py3",
-        framework_version=SKLEARN_VERSION,
+        framework_version=defaults.SKLEARN_VERSION,
         predictor_cls=SKLearnPredictor,
         model_server_workers=None,
         **kwargs
@@ -108,7 +108,9 @@ def __init__(
         )
 
         if py_version == "py2":
-            logger.warning(python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION))
+            logger.warning(
+                python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
+            )
 
         self.py_version = py_version
         self.framework_version = framework_version

src/sagemaker/tensorflow/estimator.py

@@ -25,7 +25,7 @@
 from sagemaker.debugger import DebuggerHookConfig
 from sagemaker.estimator import Framework
 import sagemaker.fw_utils as fw
-from sagemaker.tensorflow.defaults import TF_VERSION, LATEST_VERSION, LATEST_PY2_VERSION
+from sagemaker.tensorflow import defaults
 from sagemaker.tensorflow.model import TensorFlowModel
 from sagemaker.tensorflow.serving import Model
 from sagemaker.transformer import Transformer
@@ -197,7 +197,7 @@ class TensorFlow(Framework):
 
     __framework_name__ = "tensorflow"
 
-    LATEST_VERSION = LATEST_VERSION
+    LATEST_VERSION = defaults.LATEST_VERSION
 
     _LATEST_1X_VERSION = "1.15.0"
 
@@ -288,14 +288,16 @@ def __init__(
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
         if framework_version is None:
-            logger.warning(fw.empty_framework_version_warning(TF_VERSION, self.LATEST_VERSION))
-        self.framework_version = framework_version or TF_VERSION
+            logger.warning(
+                fw.empty_framework_version_warning(defaults.TF_VERSION, self.LATEST_VERSION)
+            )
+        self.framework_version = framework_version or defaults.TF_VERSION
 
         if not py_version:
             py_version = "py3" if self._only_python_3_supported() else "py2"
         if py_version == "py2":
             logger.warning(
-                fw.python_deprecation_warning(self.__framework_name__, LATEST_PY2_VERSION)
+                fw.python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
             )
 
         if "enable_sagemaker_metrics" not in kwargs:
@@ -360,8 +362,8 @@ def _validate_args(
 
         if py_version == "py2" and self._only_python_3_supported():
             msg = (
-                "Python 2 containers are only available before January 1st, 2020. "
-                "Please use a Python 3 container."
+                "Python 2 containers are only available with {} and lower versions. "
+                "Please use a Python 3 container.".format(defaults.LATEST_PY2_VERSION)
             )
             raise AttributeError(msg)