breaking: require framework_version, py_version for tensorflow

metrizable · metrizable · commit 99279ffb7318 · 2020-06-11T16:38:09.000-07:00
diff --git a/doc/frameworks/tensorflow/upgrade_from_legacy.rst b/doc/frameworks/tensorflow/upgrade_from_legacy.rst
@@ -104,6 +104,7 @@ the difference in code would be as follows:
         ...
         source_dir="code",
         framework_version="1.10.0",
+        py_version="py3",
         train_instance_type="ml.m4.xlarge",
         image_name="520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow:1.10.0-cpu-py2",
         hyperparameters={
diff --git a/doc/frameworks/tensorflow/using_tf.rst b/doc/frameworks/tensorflow/using_tf.rst
@@ -319,7 +319,7 @@ To run training job with Pipe input mode, pass in ``input_mode='Pipe'`` to your
     tf_estimator = TensorFlow(entry_point='tf-train-with-pipemodedataset.py', role='SageMakerRole',
                               training_steps=10000, evaluation_steps=100,
                               train_instance_count=1, train_instance_type='ml.p2.xlarge',
-                              framework_version='1.10.0', input_mode='Pipe')
+                              framework_version='1.10.0', py_version='py3', input_mode='Pipe')
 
     tf_estimator.fit('s3://bucket/path/to/training/data')
 
@@ -383,7 +383,8 @@ estimator object to create a SageMaker Endpoint:
   from sagemaker.tensorflow import TensorFlow
 
   estimator = TensorFlow(entry_point='tf-train.py', ..., train_instance_count=1,
-                         train_instance_type='ml.c4.xlarge', framework_version='1.11')
+                         train_instance_type='ml.c4.xlarge', framework_version='1.11',
+                         py_version='py3')
 
   estimator.fit(inputs)
 
diff --git a/src/sagemaker/rl/estimator.py b/src/sagemaker/rl/estimator.py
@@ -254,7 +254,9 @@ def create_model(
             )
 
         if self.framework == RLFramework.TENSORFLOW.value:
-            return TensorFlowModel(framework_version=self.framework_version, **base_args)
+            return TensorFlowModel(
+                framework_version=self.framework_version, py_version=PYTHON_VERSION, **base_args
+            )
         if self.framework == RLFramework.MXNET.value:
             return MXNetModel(
                 framework_version=self.framework_version, py_version=PYTHON_VERSION, **extended_args
diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py
@@ -56,11 +56,12 @@ def __init__(
 
         Args:
             py_version (str): Python version you want to use for executing your model training
-                code (default: 'py2').
+                code. One of 'py2', 'py3', or 'py37'. Defaults to ``None``. Required unless
+                ``image_name`` is provided.
             framework_version (str): TensorFlow version you want to use for executing your model
-                training code. List of supported versions
+                training code. Defaults to ``None``. Required unless ``image_name`` is provided.
+                List of supported versions:
                 https://github.com/aws/sagemaker-python-sdk#tensorflow-sagemaker-estimators.
-                If not specified, this will default to 1.11.
             model_dir (str): S3 location where the checkpoint data and models can be exported to
                 during training (default: None). It will be passed in the training script as one of
                 the command line arguments. If not specified, one is provided based on
@@ -81,6 +82,10 @@ def __init__(
                 Examples:
                     123.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:1.0
                     custom-image:latest.
+
+                If ``framework_version`` or ``py_version`` are ``None``, then
+                ``image_name`` is required. If also ``None``, then a ``ValueError``
+                will be raised.
             distributions (dict): A dictionary with information on how to run distributed training
                 (default: None). Currently we support distributed training with parameter servers
                 and MPI.
@@ -114,18 +119,13 @@ def __init__(
             :class:`~sagemaker.estimator.Framework` and
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
-        if framework_version is None:
-            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"
+        fw.validate_version_or_image_args(framework_version, py_version, image_name)
         if py_version == "py2":
             logger.warning(
                 fw.python_deprecation_warning(self.__framework_name__, defaults.LATEST_PY2_VERSION)
             )
+        self.framework_version = framework_version
+        self.py_version = py_version
 
         if distributions is not None:
             logger.warning(fw.parameter_v2_rename_warning("distribution", distributions))
@@ -136,12 +136,11 @@ def __init__(
 
         if "enable_sagemaker_metrics" not in kwargs:
             # enable sagemaker metrics for TF v1.15 or greater:
-            if fw.is_version_equal_or_higher([1, 15], self.framework_version):
+            if framework_version and fw.is_version_equal_or_higher([1, 15], framework_version):
                 kwargs["enable_sagemaker_metrics"] = True
 
         super(TensorFlow, self).__init__(image_name=image_name, **kwargs)
 
-        self.py_version = py_version
         self.model_dir = model_dir
         self.distributions = distributions or {}
 
@@ -150,7 +149,7 @@ def __init__(
     def _validate_args(self, py_version, framework_version):
         """Placeholder docstring"""
 
-        if py_version == "py3":
+        if py_version:
             if framework_version is None:
                 raise AttributeError(fw.EMPTY_FRAMEWORK_VERSION_ERROR)
 
@@ -161,7 +160,7 @@ def _validate_args(self, py_version, framework_version):
             )
             raise AttributeError(msg)
 
-        if self._only_legacy_mode_supported() and self.image_name is None:
+        if self.image_name is None and self._only_legacy_mode_supported():
             legacy_image_uri = fw.create_image_uri(
                 self.sagemaker_session.boto_region_name,
                 "tensorflow",
@@ -294,6 +293,7 @@ def create_model(
             role=role or self.role,
             container_log_level=self.container_log_level,
             framework_version=self.framework_version,
+            py_version=self.py_version,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
             entry_point=entry_point,
diff --git a/src/sagemaker/tensorflow/model.py b/src/sagemaker/tensorflow/model.py
@@ -17,9 +17,8 @@
 
 import sagemaker
 from sagemaker.content_types import CONTENT_TYPE_JSON
-from sagemaker.fw_utils import create_image_uri
+from sagemaker.fw_utils import create_image_uri, validate_version_or_image_args
 from sagemaker.predictor import json_serializer, json_deserializer
-from sagemaker.tensorflow.defaults import TF_VERSION
 
 
 class TensorFlowPredictor(sagemaker.RealTimePredictor):
@@ -138,7 +137,8 @@ def __init__(
         role,
         entry_point=None,
         image=None,
-        framework_version=TF_VERSION,
+        framework_version=None,
+        py_version=None,
         container_log_level=None,
         predictor_cls=TensorFlowPredictor,
         **kwargs
@@ -158,9 +158,16 @@ def __init__(
                 hosting. If ``source_dir`` is specified, then ``entry_point``
                 must point to a file located at the root of ``source_dir``.
             image (str): A Docker image URI (default: None). If not specified, a
-                default image for TensorFlow Serving will be used.
+                default image for TensorFlow Serving will be used. If
+                ``framework_version`` or ``py_version`` are ``None``, then
+                ``image`` is required. If also ``None``, then a ``ValueError``
+                will be raised.
             framework_version (str): Optional. TensorFlow Serving version you
-                want to use.
+                want to use. Defaults to ``None``. Required unless ``image`` is
+                provided.
+            py_version (str): Python version you want to use for executing your
+                model training code. One of 'py2', 'py3', or 'py37'. Defaults to
+                ``None``. Required unless ``image`` is provided.
             container_log_level (int): Log level to use within the container
                 (default: logging.ERROR). Valid values are defined in the Python
                 logging module.
@@ -176,6 +183,10 @@ def __init__(
             :class:`~sagemaker.model.FrameworkModel` and
             :class:`~sagemaker.model.Model`.
         """
+        validate_version_or_image_args(framework_version, py_version, image)
+        self.framework_version = framework_version
+        self.py_version = py_version
+
         super(TensorFlowModel, self).__init__(
             model_data=model_data,
             role=role,
@@ -184,7 +195,6 @@ def __init__(
             entry_point=entry_point,
             **kwargs
         )
-        self.framework_version = framework_version
         self._container_log_level = container_log_level
 
     def deploy(
diff --git a/tests/integ/test_data_capture_config.py b/tests/integ/test_data_capture_config.py
@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import os
+import pytest
 
 import sagemaker
 import tests.integ
@@ -40,8 +41,13 @@
 CUSTOM_JSON_CONTENT_TYPES = ["application/jsontype1", "application/jsontype2"]
 
 
+@pytest.fixture(scope="module")
+def py_version(tf_full_version, tf_serving_version):
+    return "py37" if tf_full_version == tf_serving_version else tests.integ.PYTHON_VERSION
+
+
 def test_enabling_data_capture_on_endpoint_shows_correct_data_capture_status(
-    sagemaker_session, tf_serving_version
+    sagemaker_session, tf_serving_version, py_version
 ):
     endpoint_name = unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
@@ -53,6 +59,7 @@ def test_enabling_data_capture_on_endpoint_shows_correct_data_capture_status(
             model_data=model_data,
             role=ROLE,
             framework_version=tf_serving_version,
+            py_version=py_version,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(
@@ -98,7 +105,7 @@ def test_enabling_data_capture_on_endpoint_shows_correct_data_capture_status(
 
 
 def test_disabling_data_capture_on_endpoint_shows_correct_data_capture_status(
-    sagemaker_session, tf_serving_version
+    sagemaker_session, tf_serving_version, py_version
 ):
     endpoint_name = unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
@@ -110,6 +117,7 @@ def test_disabling_data_capture_on_endpoint_shows_correct_data_capture_status(
             model_data=model_data,
             role=ROLE,
             framework_version=tf_serving_version,
+            py_version=py_version,
             sagemaker_session=sagemaker_session,
         )
         destination_s3_uri = os.path.join(
@@ -184,7 +192,7 @@ def test_disabling_data_capture_on_endpoint_shows_correct_data_capture_status(
 
 
 def test_updating_data_capture_on_endpoint_shows_correct_data_capture_status(
-    sagemaker_session, tf_serving_version
+    sagemaker_session, tf_serving_version, py_version
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
@@ -196,6 +204,7 @@ def test_updating_data_capture_on_endpoint_shows_correct_data_capture_status(
             model_data=model_data,
             role=ROLE,
             framework_version=tf_serving_version,
+            py_version=py_version,
             sagemaker_session=sagemaker_session,
         )
         destination_s3_uri = os.path.join(
diff --git a/tests/integ/test_model_monitor.py b/tests/integ/test_model_monitor.py
@@ -88,7 +88,12 @@
 
 
 @pytest.fixture(scope="module")
-def predictor(sagemaker_session, tf_serving_version):
+def py_version(tf_full_version, tf_serving_version):
+    return "py37" if tf_full_version == tf_serving_version else tests.integ.PYTHON_VERSION
+
+
+@pytest.fixture(scope="module")
+def predictor(sagemaker_session, tf_serving_version, py_version):
     endpoint_name = unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
         path=os.path.join(tests.integ.DATA_DIR, "tensorflow-serving-test-model.tar.gz"),
@@ -101,6 +106,7 @@ def predictor(sagemaker_session, tf_serving_version):
             model_data=model_data,
             role=ROLE,
             framework_version=tf_serving_version,
+            py_version=py_version,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(
diff --git a/tests/integ/test_tf.py b/tests/integ/test_tf.py
@@ -59,7 +59,7 @@ def test_mnist_with_checkpoint_config(
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,
         framework_version=tf_full_version,
-        py_version="py37",
+        py_version=py_version,
         metric_definitions=[{"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}],
         checkpoint_s3_uri=checkpoint_s3_uri,
         checkpoint_local_path=checkpoint_local_path,
@@ -137,8 +137,8 @@ def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version, py
         train_instance_count=2,
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,
-        py_version="py37",
         framework_version=tf_full_version,
+        py_version=py_version,
         distributions=PARAMETER_SERVER_DISTRIBUTION,
     )
     inputs = estimator.sagemaker_session.upload_data(
diff --git a/tests/integ/test_tf_efs_fsx.py b/tests/integ/test_tf_efs_fsx.py
@@ -36,7 +36,7 @@
 FSX_DIR_PATH = "/fsx/tensorflow"
 MAX_JOBS = 2
 MAX_PARALLEL_JOBS = 2
-PY_VERSION = "py3"
+PY_VERSION = "py37"
 
 
 @pytest.fixture(scope="module")
@@ -139,8 +139,8 @@ def test_tuning_tf_efs(efs_fsx_setup, sagemaker_session, cpu_instance_type):
         train_instance_count=1,
         train_instance_type=cpu_instance_type,
         sagemaker_session=sagemaker_session,
-        py_version=PY_VERSION,
         framework_version=TensorFlow.LATEST_VERSION,
+        py_version=PY_VERSION,
         subnets=subnets,
         security_group_ids=security_group_ids,
     )
@@ -186,8 +186,8 @@ def test_tuning_tf_lustre(efs_fsx_setup, sagemaker_session, cpu_instance_type):
         train_instance_count=1,
         train_instance_type=cpu_instance_type,
         sagemaker_session=sagemaker_session,
-        py_version=PY_VERSION,
         framework_version=TensorFlow.LATEST_VERSION,
+        py_version=PY_VERSION,
         subnets=subnets,
         security_group_ids=security_group_ids,
     )
diff --git a/tests/integ/test_tfs.py b/tests/integ/test_tfs.py
@@ -27,7 +27,12 @@
 
 
 @pytest.fixture(scope="module")
-def tfs_predictor(sagemaker_session, tf_serving_version):
+def py_version(tf_full_version, tf_serving_version):
+    return "py37" if tf_full_version == tf_serving_version else tests.integ.PYTHON_VERSION
+
+
+@pytest.fixture(scope="module")
+def tfs_predictor(sagemaker_session, tf_serving_version, py_version):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
         path=os.path.join(tests.integ.DATA_DIR, "tensorflow-serving-test-model.tar.gz"),
@@ -38,6 +43,7 @@ def tfs_predictor(sagemaker_session, tf_serving_version):
             model_data=model_data,
             role="SageMakerRole",
             framework_version=tf_serving_version,
+            py_version=py_version,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(1, "ml.c5.xlarge", endpoint_name=endpoint_name)
@@ -54,7 +60,7 @@ def tar_dir(directory, tmpdir):
 
 @pytest.fixture
 def tfs_predictor_with_model_and_entry_point_same_tar(
-    sagemaker_local_session, tf_serving_version, tmpdir
+    sagemaker_local_session, tf_serving_version, py_version, tmpdir
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
 
@@ -66,6 +72,7 @@ def tfs_predictor_with_model_and_entry_point_same_tar(
         model_data="file://" + model_tar,
         role="SageMakerRole",
         framework_version=tf_serving_version,
+        py_version=py_version,
         sagemaker_session=sagemaker_local_session,
     )
     predictor = model.deploy(1, "local", endpoint_name=endpoint_name)
@@ -78,7 +85,7 @@ def tfs_predictor_with_model_and_entry_point_same_tar(
 
 @pytest.fixture(scope="module")
 def tfs_predictor_with_model_and_entry_point_and_dependencies(
-    sagemaker_local_session, tf_serving_version
+    sagemaker_local_session, tf_serving_version, py_version
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
 
@@ -99,6 +106,7 @@ def tfs_predictor_with_model_and_entry_point_and_dependencies(
         role="SageMakerRole",
         dependencies=dependencies,
         framework_version=tf_serving_version,
+        py_version=py_version,
         sagemaker_session=sagemaker_local_session,
     )
 
@@ -122,6 +130,7 @@ def tfs_predictor_with_accelerator(sagemaker_session, ei_tf_full_version, cpu_in
             model_data=model_data,
             role="SageMakerRole",
             framework_version=ei_tf_full_version,
+            py_version=tests.integ.PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(
diff --git a/tests/integ/test_tuner.py b/tests/integ/test_tuner.py
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
diff --git a/tests/unit/sagemaker/tensorflow/test_estimator.py b/tests/unit/sagemaker/tensorflow/test_estimator.py
diff --git a/tests/unit/sagemaker/tensorflow/test_estimator_init.py b/tests/unit/sagemaker/tensorflow/test_estimator_init.py
diff --git a/tests/unit/sagemaker/tensorflow/test_tfs.py b/tests/unit/sagemaker/tensorflow/test_tfs.py
