fix: updates based on PR feedback

Author: metrizable

doc/frameworks/tensorflow/upgrade_from_legacy.rst

@@ -104,7 +104,7 @@ the difference in code would be as follows:
         ...
         source_dir="code",
         framework_version="1.10.0",
-        py_version="py3",
+        py_version="py2",
         train_instance_type="ml.m4.xlarge",
         image_name="520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow:1.10.0-cpu-py2",
         hyperparameters={

src/sagemaker/rl/estimator.py

@@ -254,9 +254,7 @@ def create_model(
             )
 
         if self.framework == RLFramework.TENSORFLOW.value:
-            return TensorFlowModel(
-                framework_version=self.framework_version, py_version=PYTHON_VERSION, **base_args
-            )
+            return TensorFlowModel(framework_version=self.framework_version, **base_args)
         if self.framework == RLFramework.MXNET.value:
             return MXNetModel(
                 framework_version=self.framework_version, py_version=PYTHON_VERSION, **extended_args

src/sagemaker/tensorflow/estimator.py

@@ -56,8 +56,7 @@ def __init__(
 
         Args:
             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_name`` is provided.
+                code. Defaults to ``None``. Required unless ``image_name`` is provided.
             framework_version (str): TensorFlow version you want to use for executing your model
                 training code. Defaults to ``None``. Required unless ``image_name`` is provided.
                 List of supported versions:
@@ -144,15 +143,11 @@ def __init__(
         self.model_dir = model_dir
         self.distributions = distributions or {}
 
-        self._validate_args(py_version=py_version, framework_version=self.framework_version)
+        self._validate_args(py_version=py_version)
 
-    def _validate_args(self, py_version, framework_version):
+    def _validate_args(self, py_version):
         """Placeholder docstring"""
 
-        if py_version:
-            if framework_version is None:
-                raise AttributeError(fw.EMPTY_FRAMEWORK_VERSION_ERROR)
-
         if py_version == "py2" and self._only_python_3_supported():
             msg = (
                 "Python 2 containers are only available with {} and lower versions. "
@@ -293,7 +288,6 @@ 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,

src/sagemaker/tensorflow/model.py

@@ -17,7 +17,7 @@
 
 import sagemaker
 from sagemaker.content_types import CONTENT_TYPE_JSON
-from sagemaker.fw_utils import create_image_uri, validate_version_or_image_args
+from sagemaker.fw_utils import create_image_uri
 from sagemaker.predictor import json_serializer, json_deserializer
 
 
@@ -138,7 +138,6 @@ def __init__(
         entry_point=None,
         image=None,
         framework_version=None,
-        py_version=None,
         container_log_level=None,
         predictor_cls=TensorFlowPredictor,
         **kwargs
@@ -159,15 +158,11 @@ def __init__(
                 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. If
-                ``framework_version`` or ``py_version`` are ``None``, then
-                ``image`` is required. If also ``None``, then a ``ValueError``
-                will be raised.
+                ``framework_version`` is ``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. 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.
@@ -183,9 +178,12 @@ def __init__(
             :class:`~sagemaker.model.FrameworkModel` and
             :class:`~sagemaker.model.Model`.
         """
-        validate_version_or_image_args(framework_version, py_version, image)
+        if framework_version is None and image is None:
+            raise ValueError(
+                "Both framework_version and image were None. "
+                "Either specify framework_version or specify image_name."
+            )
         self.framework_version = framework_version
-        self.py_version = py_version
 
         super(TensorFlowModel, self).__init__(
             model_data=model_data,

tests/conftest.py

@@ -209,7 +209,6 @@ def xgboost_version(request):
         "1.11.0",
         "1.12",
         "1.12.0",
-        "1.13",
         "1.14",
         "1.14.0",
         "1.15",
@@ -226,6 +225,16 @@ def tf_version(request):
     return request.param
 
 
+@pytest.fixture(scope="module", params=["py2", "py3"])
+def tf_py_version(tf_version, request):
+    version = [int(val) for val in tf_version.split(".")]
+    if version < [1, 13]:
+        return "py2"
+    if version < [2, 2]:
+        return request.param
+    return "py37"
+
+
 @pytest.fixture(scope="module", params=["0.10.1", "0.10.1", "0.11", "0.11.0", "0.11.1"])
 def rl_coach_tf_version(request):
     return request.param
@@ -290,6 +299,23 @@ def tf_full_version(request):
         return tf_version
 
 
+@pytest.fixture(scope="module")
+def tf_full_py_version(tf_full_version, request):
+    """fixture to match tf_full_version
+
+    Fixture exists as such, since tf_full_version may be overridden --tf-full-version.
+    Otherwise, this would simply be py37 to match the latest version support.
+
+    TODO: Evaluate use of --tf-full-version with possible eye to remove and simplify code.
+    """
+    version = [int(val) for val in tf_full_version.split(".")]
+    if version < [1, 13]:
+        return "py2"
+    if tf_full_version in [TensorFlow._LATEST_1X_VERSION, LATEST_VERSION]:
+        return "py37"
+    return "py3"
+
+
 @pytest.fixture(scope="module", params=["1.15.0", "2.0.0"])
 def ei_tf_full_version(request):
     tf_ei_version = request.config.getoption("--ei-tf-full-version")

tests/integ/test_airflow_config.py

@@ -562,7 +562,7 @@ def test_tf_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_inst
             train_instance_type=cpu_instance_type,
             sagemaker_session=sagemaker_session,
             framework_version=TensorFlow.LATEST_VERSION,
-            py_version=PYTHON_VERSION,
+            py_version="py37",  # only version available with 2.2.0
             metric_definitions=[
                 {"Name": "train:global_steps", "Regex": r"global_step\/sec:\s(.*)"}
             ],

tests/integ/test_data_capture_config.py

@@ -20,7 +20,6 @@
 from sagemaker.model_monitor import DataCaptureConfig, NetworkConfig
 from sagemaker.tensorflow.model import TensorFlowModel
 from sagemaker.utils import unique_name_from_base
-from tests.integ import PYTHON_VERSION
 from tests.integ.retry import retries
 
 ROLE = "SageMakerRole"
@@ -54,7 +53,6 @@ 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=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(
@@ -112,7 +110,6 @@ 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=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         destination_s3_uri = os.path.join(
@@ -199,7 +196,6 @@ 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=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         destination_s3_uri = os.path.join(

tests/integ/test_model_monitor.py

@@ -24,7 +24,7 @@
 from sagemaker.s3 import S3Uploader
 from datetime import datetime, timedelta
 
-from tests.integ import DATA_DIR, PYTHON_VERSION
+from tests.integ import DATA_DIR
 from sagemaker.model_monitor import DatasetFormat
 from sagemaker.model_monitor import NetworkConfig, Statistics, Constraints
 from sagemaker.model_monitor import ModelMonitor
@@ -101,7 +101,6 @@ def predictor(sagemaker_session, tf_serving_version):
             model_data=model_data,
             role=ROLE,
             framework_version=tf_serving_version,
-            py_version=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(

tests/integ/test_tf.py

@@ -19,7 +19,7 @@
 import pytest
 
 from sagemaker.tensorflow import TensorFlow
-from sagemaker.tensorflow.defaults import LATEST_VERSION, LATEST_SERVING_VERSION
+from sagemaker.tensorflow.defaults import LATEST_SERVING_VERSION
 from sagemaker.utils import unique_name_from_base, sagemaker_timestamp
 
 import tests.integ
@@ -38,16 +38,9 @@
 MPI_DISTRIBUTION = {"mpi": {"enabled": True}}
 TAGS = [{"Key": "some-key", "Value": "some-value"}]
 
-PY37_SUPPORTED_FRAMEWORK_VERSION = [TensorFlow._LATEST_1X_VERSION, LATEST_VERSION]
-
-
-@pytest.fixture(scope="module")
-def py_version(tf_full_version):
-    return "py37" if tf_full_version in PY37_SUPPORTED_FRAMEWORK_VERSION else PYTHON_VERSION
-
 
 def test_mnist_with_checkpoint_config(
-    sagemaker_session, instance_type, tf_full_version, py_version
+    sagemaker_session, instance_type, tf_full_version, tf_full_py_version
 ):
     checkpoint_s3_uri = "s3://{}/checkpoints/tf-{}".format(
         sagemaker_session.default_bucket(), sagemaker_timestamp()
@@ -60,7 +53,7 @@ def test_mnist_with_checkpoint_config(
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,
         framework_version=tf_full_version,
-        py_version=py_version,
+        py_version=tf_full_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,
@@ -131,15 +124,15 @@ def test_server_side_encryption(sagemaker_session, tf_serving_version):
 
 
 @pytest.mark.canary_quick
-def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version, py_version):
+def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version, tf_full_py_version):
     estimator = TensorFlow(
         entry_point=SCRIPT,
         role=ROLE,
         train_instance_count=2,
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,
         framework_version=tf_full_version,
-        py_version=py_version,
+        py_version=tf_full_py_version,
         distributions=PARAMETER_SERVER_DISTRIBUTION,
     )
     inputs = estimator.sagemaker_session.upload_data(

tests/integ/test_tfs.py

@@ -24,7 +24,6 @@
 import tests.integ
 import tests.integ.timeout
 from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor
-from tests.integ import PYTHON_VERSION
 
 
 @pytest.fixture(scope="module")
@@ -39,7 +38,6 @@ def tfs_predictor(sagemaker_session, tf_serving_version):
             model_data=model_data,
             role="SageMakerRole",
             framework_version=tf_serving_version,
-            py_version=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(1, "ml.c5.xlarge", endpoint_name=endpoint_name)
@@ -68,7 +66,6 @@ def tfs_predictor_with_model_and_entry_point_same_tar(
         model_data="file://" + model_tar,
         role="SageMakerRole",
         framework_version=tf_serving_version,
-        py_version=PYTHON_VERSION,
         sagemaker_session=sagemaker_local_session,
     )
     predictor = model.deploy(1, "local", endpoint_name=endpoint_name)
@@ -102,7 +99,6 @@ def tfs_predictor_with_model_and_entry_point_and_dependencies(
         role="SageMakerRole",
         dependencies=dependencies,
         framework_version=tf_serving_version,
-        py_version=PYTHON_VERSION,
         sagemaker_session=sagemaker_local_session,
     )
 
@@ -126,7 +122,6 @@ 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=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(

tests/integ/test_tuner.py

@@ -37,7 +37,6 @@
 from sagemaker.predictor import json_deserializer
 from sagemaker.pytorch import PyTorch
 from sagemaker.tensorflow import TensorFlow
-from sagemaker.tensorflow.defaults import LATEST_VERSION
 from sagemaker.tuner import (
     IntegerParameter,
     ContinuousParameter,
@@ -52,13 +51,6 @@
 
 DATA_PATH = os.path.join(DATA_DIR, "iris", "data")
 
-PY37_SUPPORTED_FRAMEWORK_VERSION = [TensorFlow._LATEST_1X_VERSION, LATEST_VERSION]
-
-
-@pytest.fixture(scope="module")
-def py_version(tf_full_version):
-    return "py37" if tf_full_version in PY37_SUPPORTED_FRAMEWORK_VERSION else PYTHON_VERSION
-
 
 @pytest.fixture(scope="module")
 def kmeans_train_set(sagemaker_session):
@@ -598,7 +590,9 @@ def test_tuning_mxnet(sagemaker_session, mxnet_full_version, cpu_instance_type):
 
 
 @pytest.mark.canary_quick
-def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type, tf_full_version, py_version):
+def test_tuning_tf_script_mode(
+    sagemaker_session, cpu_instance_type, tf_full_version, tf_full_py_version
+):
     resource_path = os.path.join(DATA_DIR, "tensorflow_mnist")
     script_path = os.path.join(resource_path, "mnist.py")
 
@@ -609,7 +603,7 @@ def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type, tf_full_ver
         train_instance_type=cpu_instance_type,
         sagemaker_session=sagemaker_session,
         framework_version=tf_full_version,
-        py_version=py_version,
+        py_version=tf_full_py_version,
     )
 
     hyperparameter_ranges = {"epochs": IntegerParameter(1, 2)}
@@ -639,9 +633,10 @@ def test_tuning_tf_script_mode(sagemaker_session, cpu_instance_type, tf_full_ver
         tuner.wait()
 
 
+# TODO: evaluate skip mark and default framework_version 1.11
 @pytest.mark.canary_quick
 @pytest.mark.skipif(PYTHON_VERSION != "py2", reason="TensorFlow image supports only python 2.")
-def test_tuning_tf(sagemaker_session, cpu_instance_type, tf_full_version, py_version):
+def test_tuning_tf(sagemaker_session, cpu_instance_type):
     with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
         script_path = os.path.join(DATA_DIR, "iris", "iris-dnn-classifier.py")
 
@@ -654,8 +649,8 @@ def test_tuning_tf(sagemaker_session, cpu_instance_type, tf_full_version, py_ver
             train_instance_count=1,
             train_instance_type=cpu_instance_type,
             sagemaker_session=sagemaker_session,
-            framework_version=tf_full_version,
-            py_version=py_version,
+            framework_version="1.11",
+            py_version=PYTHON_VERSION,
         )
 
         inputs = sagemaker_session.upload_data(path=DATA_PATH, key_prefix="integ-test-data/tf_iris")
@@ -695,8 +690,9 @@ def test_tuning_tf(sagemaker_session, cpu_instance_type, tf_full_version, py_ver
         assert dict_result == list_result
 
 
+# TODO: evaluate skip mark and default framework_version 1.11
 @pytest.mark.skipif(PYTHON_VERSION != "py2", reason="TensorFlow image supports only python 2.")
-def test_tuning_tf_vpc_multi(sagemaker_session, cpu_instance_type, tf_full_version, py_version):
+def test_tuning_tf_vpc_multi(sagemaker_session, cpu_instance_type):
     """Test Tensorflow multi-instance using the same VpcConfig for training and inference"""
     instance_type = cpu_instance_type
     instance_count = 2
@@ -720,8 +716,8 @@ def test_tuning_tf_vpc_multi(sagemaker_session, cpu_instance_type, tf_full_versi
         subnets=subnet_ids,
         security_group_ids=[security_group_id],
         encrypt_inter_container_traffic=True,
-        framework_version=tf_full_version,
-        py_version=py_version,
+        framework_version="1.11",
+        py_version=PYTHON_VERSION,
     )
 
     inputs = sagemaker_session.upload_data(path=DATA_PATH, key_prefix="integ-test-data/tf_iris")