fix: add default framework version warning message in Model classes (#1218)

chuyang-deng · web-flow · commit 926690e216dc · 2020-01-09T16:41:23.000-08:00
* fix: add default framework version warning message in Model classes

* fix black-format errors

* add LATEST_VERSION in frameowrk defaults.py
diff --git a/src/sagemaker/chainer/defaults.py b/src/sagemaker/chainer/defaults.py
@@ -17,3 +17,6 @@
 """Default Chainer version for when the framework version is not specified.
 This is no longer updated so as to not break existing workflows.
 """
+
+LATEST_VERSION = "5.0.0"
+"""The latest version of Chainer included in the SageMaker pre-built Docker images."""
diff --git a/src/sagemaker/chainer/estimator.py b/src/sagemaker/chainer/estimator.py
@@ -22,7 +22,7 @@
     empty_framework_version_warning,
     python_deprecation_warning,
 )
-from sagemaker.chainer.defaults import CHAINER_VERSION
+from sagemaker.chainer.defaults import CHAINER_VERSION, LATEST_VERSION
 from sagemaker.chainer.model import ChainerModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -40,8 +40,7 @@ class Chainer(Framework):
     _process_slots_per_host = "sagemaker_process_slots_per_host"
     _additional_mpi_options = "sagemaker_additional_mpi_options"
 
-    LATEST_VERSION = "5.0.0"
-    """The latest version of Chainer included in the SageMaker pre-built Docker images."""
+    LATEST_VERSION = LATEST_VERSION
 
     def __init__(
         self,
diff --git a/src/sagemaker/chainer/model.py b/src/sagemaker/chainer/model.py
@@ -16,9 +16,14 @@
 import logging
 
 import sagemaker
-from sagemaker.fw_utils import create_image_uri, model_code_key_prefix, python_deprecation_warning
+from sagemaker.fw_utils import (
+    create_image_uri,
+    model_code_key_prefix,
+    python_deprecation_warning,
+    empty_framework_version_warning,
+)
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
-from sagemaker.chainer.defaults import CHAINER_VERSION
+from sagemaker.chainer.defaults import CHAINER_VERSION, LATEST_VERSION
 from sagemaker.predictor import RealTimePredictor, npy_serializer, numpy_deserializer
 
 logger = logging.getLogger("sagemaker")
@@ -61,7 +66,7 @@ def __init__(
         entry_point,
         image=None,
         py_version="py3",
-        framework_version=CHAINER_VERSION,
+        framework_version=None,
         predictor_cls=ChainerPredictor,
         model_server_workers=None,
         **kwargs
@@ -107,9 +112,11 @@ def __init__(
         )
         if py_version == "py2":
             logger.warning(python_deprecation_warning(self.__framework_name__))
+        if framework_version is None:
+            logger.warning(empty_framework_version_warning(CHAINER_VERSION, LATEST_VERSION))
 
         self.py_version = py_version
-        self.framework_version = framework_version
+        self.framework_version = framework_version or CHAINER_VERSION
         self.model_server_workers = model_server_workers
 
     def prepare_container_def(self, instance_type, accelerator_type=None):
diff --git a/src/sagemaker/mxnet/defaults.py b/src/sagemaker/mxnet/defaults.py
@@ -17,3 +17,6 @@
 """Default MXNet version for when the framework version is not specified.
 This is no longer updated so as to not break existing workflows.
 """
+
+LATEST_VERSION = "1.6.0"
+"""The latest version of MXNet included in the SageMaker pre-built Docker images."""
diff --git a/src/sagemaker/mxnet/estimator.py b/src/sagemaker/mxnet/estimator.py
@@ -23,7 +23,7 @@
     python_deprecation_warning,
     is_version_equal_or_higher,
 )
-from sagemaker.mxnet.defaults import MXNET_VERSION
+from sagemaker.mxnet.defaults import MXNET_VERSION, LATEST_VERSION
 from sagemaker.mxnet.model import MXNetModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -36,8 +36,7 @@ class MXNet(Framework):
     __framework_name__ = "mxnet"
     _LOWEST_SCRIPT_MODE_VERSION = ["1", "3"]
 
-    LATEST_VERSION = "1.6.0"
-    """The latest version of MXNet included in the SageMaker pre-built Docker images."""
+    LATEST_VERSION = LATEST_VERSION
 
     def __init__(
         self,
diff --git a/src/sagemaker/mxnet/model.py b/src/sagemaker/mxnet/model.py
@@ -18,9 +18,14 @@
 from pkg_resources import parse_version
 
 import sagemaker
-from sagemaker.fw_utils import create_image_uri, model_code_key_prefix, python_deprecation_warning
+from sagemaker.fw_utils import (
+    create_image_uri,
+    model_code_key_prefix,
+    python_deprecation_warning,
+    empty_framework_version_warning,
+)
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
-from sagemaker.mxnet.defaults import MXNET_VERSION
+from sagemaker.mxnet.defaults import MXNET_VERSION, LATEST_VERSION
 from sagemaker.predictor import RealTimePredictor, json_serializer, json_deserializer
 
 logger = logging.getLogger("sagemaker")
@@ -62,7 +67,7 @@ def __init__(
         entry_point,
         image=None,
         py_version="py2",
-        framework_version=MXNET_VERSION,
+        framework_version=None,
         predictor_cls=MXNetPredictor,
         model_server_workers=None,
         **kwargs
@@ -109,9 +114,11 @@ def __init__(
 
         if py_version == "py2":
             logger.warning(python_deprecation_warning(self.__framework_name__))
+        if framework_version is None:
+            logger.warning(empty_framework_version_warning(MXNET_VERSION, LATEST_VERSION))
 
         self.py_version = py_version
-        self.framework_version = framework_version
+        self.framework_version = framework_version or MXNET_VERSION
         self.model_server_workers = model_server_workers
 
     def prepare_container_def(self, instance_type, accelerator_type=None):
diff --git a/src/sagemaker/pytorch/defaults.py b/src/sagemaker/pytorch/defaults.py
@@ -19,4 +19,7 @@
 break existing workflows.
 """
 
+LATEST_VERSION = "1.3.1"
+"""The latest version of PyTorch included in the SageMaker pre-built Docker images."""
+
 PYTHON_VERSION = "py3"
diff --git a/src/sagemaker/pytorch/estimator.py b/src/sagemaker/pytorch/estimator.py
@@ -23,7 +23,7 @@
     python_deprecation_warning,
     is_version_equal_or_higher,
 )
-from sagemaker.pytorch.defaults import PYTORCH_VERSION, PYTHON_VERSION
+from sagemaker.pytorch.defaults import PYTORCH_VERSION, PYTHON_VERSION, LATEST_VERSION
 from sagemaker.pytorch.model import PyTorchModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
@@ -35,8 +35,7 @@ class PyTorch(Framework):
 
     __framework_name__ = "pytorch"
 
-    LATEST_VERSION = "1.3.1"
-    """The latest version of PyTorch included in the SageMaker pre-built Docker images."""
+    LATEST_VERSION = LATEST_VERSION
 
     def __init__(
         self,
diff --git a/src/sagemaker/pytorch/model.py b/src/sagemaker/pytorch/model.py
@@ -17,9 +17,14 @@
 import pkg_resources
 
 import sagemaker
-from sagemaker.fw_utils import create_image_uri, model_code_key_prefix, python_deprecation_warning
+from sagemaker.fw_utils import (
+    create_image_uri,
+    model_code_key_prefix,
+    python_deprecation_warning,
+    empty_framework_version_warning,
+)
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
-from sagemaker.pytorch.defaults import PYTORCH_VERSION, PYTHON_VERSION
+from sagemaker.pytorch.defaults import PYTORCH_VERSION, PYTHON_VERSION, LATEST_VERSION
 from sagemaker.predictor import RealTimePredictor, npy_serializer, numpy_deserializer
 
 logger = logging.getLogger("sagemaker")
@@ -63,7 +68,7 @@ def __init__(
         entry_point,
         image=None,
         py_version=PYTHON_VERSION,
-        framework_version=PYTORCH_VERSION,
+        framework_version=None,
         predictor_cls=PyTorchPredictor,
         model_server_workers=None,
         **kwargs
@@ -110,9 +115,11 @@ def __init__(
 
         if py_version == "py2":
             logger.warning(python_deprecation_warning(self.__framework_name__))
+        if framework_version is None:
+            logger.warning(empty_framework_version_warning(PYTORCH_VERSION, LATEST_VERSION))
 
         self.py_version = py_version
-        self.framework_version = framework_version
+        self.framework_version = framework_version or PYTORCH_VERSION
         self.model_server_workers = model_server_workers
 
     def prepare_container_def(self, instance_type, accelerator_type=None):
diff --git a/src/sagemaker/tensorflow/defaults.py b/src/sagemaker/tensorflow/defaults.py
@@ -17,3 +17,6 @@
 """Default TF version for when the framework version is not specified.
 This is no longer updated so as to not break existing workflows.
 """
+
+LATEST_VERSION = "2.0.0"
+"""The latest version of TensorFlow included in the SageMaker pre-built Docker images."""
diff --git a/src/sagemaker/tensorflow/estimator.py b/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
+from sagemaker.tensorflow.defaults import TF_VERSION, LATEST_VERSION
 from sagemaker.tensorflow.model import TensorFlowModel
 from sagemaker.tensorflow.serving import Model
 from sagemaker.transformer import Transformer
@@ -197,8 +197,7 @@ class TensorFlow(Framework):
 
     __framework_name__ = "tensorflow"
 
-    LATEST_VERSION = "2.0.0"
-    """The latest version of TensorFlow included in the SageMaker pre-built Docker images."""
+    LATEST_VERSION = LATEST_VERSION
 
     _LATEST_1X_VERSION = "1.15.0"
 
diff --git a/src/sagemaker/tensorflow/model.py b/src/sagemaker/tensorflow/model.py
@@ -16,10 +16,15 @@
 import logging
 
 import sagemaker
-from sagemaker.fw_utils import create_image_uri, model_code_key_prefix, python_deprecation_warning
+from sagemaker.fw_utils import (
+    create_image_uri,
+    model_code_key_prefix,
+    python_deprecation_warning,
+    empty_framework_version_warning,
+)
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.predictor import RealTimePredictor
-from sagemaker.tensorflow.defaults import TF_VERSION
+from sagemaker.tensorflow.defaults import TF_VERSION, LATEST_VERSION
 from sagemaker.tensorflow.predictor import tf_json_serializer, tf_json_deserializer
 
 logger = logging.getLogger("sagemaker")
@@ -60,7 +65,7 @@ def __init__(
         entry_point,
         image=None,
         py_version="py2",
-        framework_version=TF_VERSION,
+        framework_version=None,
         predictor_cls=TensorFlowPredictor,
         model_server_workers=None,
         **kwargs
@@ -107,9 +112,11 @@ def __init__(
 
         if py_version == "py2":
             logger.warning(python_deprecation_warning(self.__framework_name__))
+        if framework_version is None:
+            logger.warning(empty_framework_version_warning(TF_VERSION, LATEST_VERSION))
 
         self.py_version = py_version
-        self.framework_version = framework_version
+        self.framework_version = framework_version or TF_VERSION
         self.model_server_workers = model_server_workers
 
     def prepare_container_def(self, instance_type, accelerator_type=None):
diff --git a/tests/unit/test_chainer.py b/tests/unit/test_chainer.py
@@ -601,3 +601,16 @@ def test_empty_framework_version(warning, sagemaker_session):
 
     assert estimator.framework_version == defaults.CHAINER_VERSION
     warning.assert_called_with(defaults.CHAINER_VERSION, Chainer.LATEST_VERSION)
+
+
+@patch("sagemaker.chainer.model.empty_framework_version_warning")
+def test_model_empty_framework_version(warning, sagemaker_session):
+    model = ChainerModel(
+        MODEL_DATA,
+        role=ROLE,
+        entry_point=SCRIPT_PATH,
+        sagemaker_session=sagemaker_session,
+        framework_version=None,
+    )
+    assert model.framework_version == defaults.CHAINER_VERSION
+    warning.assert_called_with(defaults.CHAINER_VERSION, defaults.LATEST_VERSION)
diff --git a/tests/unit/test_mxnet.py b/tests/unit/test_mxnet.py
@@ -722,6 +722,19 @@ def test_empty_framework_version(warning, sagemaker_session):
     warning.assert_called_with(defaults.MXNET_VERSION, mx.LATEST_VERSION)
 
 
+@patch("sagemaker.mxnet.model.empty_framework_version_warning")
+def test_model_empty_framework_version(warning, sagemaker_session):
+    model = MXNetModel(
+        MODEL_DATA,
+        role=ROLE,
+        entry_point=SCRIPT_PATH,
+        sagemaker_session=sagemaker_session,
+        framework_version=None,
+    )
+    assert model.framework_version == defaults.MXNET_VERSION
+    warning.assert_called_with(defaults.MXNET_VERSION, defaults.LATEST_VERSION)
+
+
 def test_create_model_with_custom_hosting_image(sagemaker_session):
     container_log_level = '"logging.INFO"'
     source_dir = "s3://mybucket/source"
diff --git a/tests/unit/test_pytorch.py b/tests/unit/test_pytorch.py
@@ -532,6 +532,20 @@ def test_empty_framework_version(warning, sagemaker_session):
     warning.assert_called_with(defaults.PYTORCH_VERSION, estimator.LATEST_VERSION)
 
 
+@patch("sagemaker.pytorch.model.empty_framework_version_warning")
+def test_model_empty_framework_version(warning, sagemaker_session):
+    model = PyTorchModel(
+        MODEL_DATA,
+        role=ROLE,
+        entry_point=SCRIPT_PATH,
+        sagemaker_session=sagemaker_session,
+        framework_version=None,
+    )
+
+    assert model.framework_version == defaults.PYTORCH_VERSION
+    warning.assert_called_with(defaults.PYTORCH_VERSION, defaults.LATEST_VERSION)
+
+
 def test_pt_enable_sm_metrics(sagemaker_session):
     pytorch = _pytorch_estimator(sagemaker_session, enable_sagemaker_metrics=True)
     assert pytorch.enable_sagemaker_metrics
diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py
@@ -938,6 +938,17 @@ def test_empty_framework_version(warning, sagemaker_session):
     assert estimator.framework_version == defaults.TF_VERSION
     warning.assert_called_with(defaults.TF_VERSION, estimator.LATEST_VERSION)
 
+    model = TensorFlowModel(
+        MODEL_DATA,
+        role=ROLE,
+        entry_point=SCRIPT_PATH,
+        sagemaker_session=sagemaker_session,
+        framework_version=None,
+    )
+
+    assert model.framework_version == defaults.TF_VERSION
+    warning.assert_called_with(defaults.TF_VERSION, defaults.LATEST_VERSION)
+
 
 def _deprecated_args_msg(args):
     return "{} are deprecated in script mode. Please do not set {}.".format(
