fix: py_version issues between tf_full and tf_serving

Author: metrizable

tests/integ/test_data_capture_config.py

@@ -13,14 +13,14 @@
 from __future__ import absolute_import
 
 import os
-import pytest
 
 import sagemaker
 import tests.integ
 import tests.integ.timeout
 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"
@@ -41,13 +41,8 @@
 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, py_version
+    sagemaker_session, tf_serving_version
 ):
     endpoint_name = unique_name_from_base("sagemaker-tensorflow-serving")
     model_data = sagemaker_session.upload_data(
@@ -59,7 +54,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,
+            py_version=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(
@@ -117,7 +112,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,
+            py_version=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         destination_s3_uri = os.path.join(
@@ -204,7 +199,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,
+            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
+from tests.integ import DATA_DIR, PYTHON_VERSION
 from sagemaker.model_monitor import DatasetFormat
 from sagemaker.model_monitor import NetworkConfig, Statistics, Constraints
 from sagemaker.model_monitor import ModelMonitor
@@ -88,12 +88,7 @@
 
 
 @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
-
-
-@pytest.fixture(scope="module")
-def predictor(sagemaker_session, tf_serving_version, py_version):
+def predictor(sagemaker_session, tf_serving_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"),
@@ -106,7 +101,7 @@ def predictor(sagemaker_session, tf_serving_version, py_version):
             model_data=model_data,
             role=ROLE,
             framework_version=tf_serving_version,
-            py_version=py_version,
+            py_version=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(

tests/integ/test_tf.py

@@ -19,12 +19,11 @@
 import pytest
 
 from sagemaker.tensorflow import TensorFlow
-from sagemaker.tensorflow.defaults import LATEST_SERVING_VERSION
+from sagemaker.tensorflow.defaults import LATEST_VERSION, LATEST_SERVING_VERSION
 from sagemaker.utils import unique_name_from_base, sagemaker_timestamp
 
 import tests.integ
-from tests.integ import timeout
-from tests.integ import kms_utils
+from tests.integ import kms_utils, timeout, PYTHON_VERSION
 from tests.integ.retry import retries
 from tests.integ.s3_utils import assert_s3_files_exist
 
@@ -39,10 +38,12 @@
 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, tf_serving_version):
-    return "py37" if tf_full_version == tf_serving_version else tests.integ.PYTHON_VERSION
+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(
@@ -89,7 +90,7 @@ def test_mnist_with_checkpoint_config(
     assert actual_training_checkpoint_config == expected_training_checkpoint_config
 
 
-def test_server_side_encryption(sagemaker_session, tf_serving_version, py_version):
+def test_server_side_encryption(sagemaker_session, tf_full_version, py_version):
     with kms_utils.bucket_with_encryption(sagemaker_session, ROLE) as (bucket_with_kms, kms_key):
         output_path = os.path.join(
             bucket_with_kms, "test-server-side-encryption", time.strftime("%y%m%d-%H%M")
@@ -102,7 +103,7 @@ def test_server_side_encryption(sagemaker_session, tf_serving_version, py_versio
             train_instance_count=1,
             train_instance_type="ml.c5.xlarge",
             sagemaker_session=sagemaker_session,
-            framework_version=tf_serving_version,
+            framework_version=tf_full_version,
             py_version=py_version,
             code_location=output_path,
             output_path=output_path,
@@ -154,13 +155,13 @@ def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version, py
     )
 
 
-def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, py_version):
+def test_mnist_async(sagemaker_session, cpu_instance_type):
     estimator = TensorFlow(
         entry_point=SCRIPT,
         role=ROLE,
         train_instance_count=1,
         train_instance_type="ml.c5.4xlarge",
-        py_version=tests.integ.PYTHON_VERSION,
+        py_version=PYTHON_VERSION,
         sagemaker_session=sagemaker_session,
         # testing py-sdk functionality, no need to run against all TF versions
         framework_version=LATEST_SERVING_VERSION,
@@ -195,18 +196,16 @@ def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, py_v
         _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name)
 
 
-def test_deploy_with_input_handlers(
-    sagemaker_session, instance_type, tf_serving_version, py_version
-):
+def test_deploy_with_input_handlers(sagemaker_session, instance_type, tf_serving_version):
     estimator = TensorFlow(
         entry_point="training.py",
         source_dir=TFS_RESOURCE_PATH,
         role=ROLE,
         train_instance_count=1,
         train_instance_type=instance_type,
-        py_version=py_version,
-        sagemaker_session=sagemaker_session,
         framework_version=tf_serving_version,
+        py_version=PYTHON_VERSION,
+        sagemaker_session=sagemaker_session,
         tags=TAGS,
     )
 

tests/integ/test_tfs.py

@@ -24,15 +24,11 @@
 import tests.integ
 import tests.integ.timeout
 from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor
+from tests.integ import PYTHON_VERSION
 
 
 @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
-
-
-@pytest.fixture(scope="module")
-def tfs_predictor(sagemaker_session, tf_serving_version, py_version):
+def tfs_predictor(sagemaker_session, tf_serving_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"),
@@ -43,7 +39,7 @@ def tfs_predictor(sagemaker_session, tf_serving_version, py_version):
             model_data=model_data,
             role="SageMakerRole",
             framework_version=tf_serving_version,
-            py_version=py_version,
+            py_version=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(1, "ml.c5.xlarge", endpoint_name=endpoint_name)
@@ -60,7 +56,7 @@ def tar_dir(directory, tmpdir):
 
 @pytest.fixture
 def tfs_predictor_with_model_and_entry_point_same_tar(
-    sagemaker_local_session, tf_serving_version, py_version, tmpdir
+    sagemaker_local_session, tf_serving_version, tmpdir
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
 
@@ -72,7 +68,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,
+        py_version=PYTHON_VERSION,
         sagemaker_session=sagemaker_local_session,
     )
     predictor = model.deploy(1, "local", endpoint_name=endpoint_name)
@@ -85,7 +81,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, py_version
+    sagemaker_local_session, tf_serving_version
 ):
     endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-tensorflow-serving")
 
@@ -106,7 +102,7 @@ def tfs_predictor_with_model_and_entry_point_and_dependencies(
         role="SageMakerRole",
         dependencies=dependencies,
         framework_version=tf_serving_version,
-        py_version=py_version,
+        py_version=PYTHON_VERSION,
         sagemaker_session=sagemaker_local_session,
     )
 
@@ -130,7 +126,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,
+            py_version=PYTHON_VERSION,
             sagemaker_session=sagemaker_session,
         )
         predictor = model.deploy(