add py_version to fixture

Author: Chuyang Deng

tests/conftest.py

@@ -133,6 +133,13 @@ def custom_bucket_name(boto_session):
     return "{}-{}-{}".format(CUSTOM_BUCKET_NAME_PREFIX, region, account)
 
 
+@pytest.fixture(scope="module")
+def py_version():
+    return (
+        "py37" if tf_full_version == TensorFlow._LATEST_1X_VERSION else tests.integ.PYTHON_VERSION
+    )
+
+
 @pytest.fixture(scope="module", params=["4.0", "4.0.0", "4.1", "4.1.0", "5.0", "5.0.0"])
 def chainer_version(request):
     return request.param

tests/integ/test_tf_script_mode.py

@@ -39,7 +39,9 @@
 TAGS = [{"Key": "some-key", "Value": "some-value"}]
 
 
-def test_mnist_with_checkpoint_config(sagemaker_session, instance_type, tf_full_version):
+def test_mnist_with_checkpoint_config(
+    sagemaker_session, instance_type, tf_full_version, py_version
+):
     checkpoint_s3_uri = "s3://{}/checkpoints/tf-{}".format(
         sagemaker_session.default_bucket(), sagemaker_timestamp()
     )
@@ -52,9 +54,7 @@ def test_mnist_with_checkpoint_config(sagemaker_session, instance_type, tf_full_
         sagemaker_session=sagemaker_session,
         script_mode=True,
         framework_version=tf_full_version,
-        py_version="py37"
-        if tf_full_version == TensorFlow._LATEST_1X_VERSION
-        else tests.integ.PYTHON_VERSION,
+        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,
@@ -84,7 +84,7 @@ def test_mnist_with_checkpoint_config(sagemaker_session, instance_type, tf_full_
     assert actual_training_checkpoint_config == expected_training_checkpoint_config
 
 
-def test_server_side_encryption(sagemaker_session, tf_full_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")
@@ -99,9 +99,7 @@ def test_server_side_encryption(sagemaker_session, tf_full_version):
             sagemaker_session=sagemaker_session,
             script_mode=True,
             framework_version=tf_full_version,
-            py_version="py37"
-            if tf_full_version == TensorFlow._LATEST_1X_VERSION
-            else tests.integ.PYTHON_VERSION,
+            py_version=py_version,
             code_location=output_path,
             output_path=output_path,
             model_dir="/opt/ml/model",
@@ -128,16 +126,14 @@ def test_server_side_encryption(sagemaker_session, tf_full_version):
 
 
 @pytest.mark.canary_quick
-def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version):
+def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version, py_version):
     estimator = TensorFlow(
         entry_point=SCRIPT,
         role=ROLE,
         train_instance_count=2,
         train_instance_type=instance_type,
         sagemaker_session=sagemaker_session,
-        py_version="py37"
-        if tf_full_version == TensorFlow._LATEST_1X_VERSION
-        else tests.integ.PYTHON_VERSION,
+        py_version=py_version,
         script_mode=True,
         framework_version=tf_full_version,
         distributions=PARAMETER_SERVER_DISTRIBUTION,
@@ -155,15 +151,13 @@ def test_mnist_distributed(sagemaker_session, instance_type, tf_full_version):
     )
 
 
-def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version):
+def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version, py_version):
     estimator = TensorFlow(
         entry_point=SCRIPT,
         role=ROLE,
         train_instance_count=1,
         train_instance_type="ml.c5.4xlarge",
-        py_version="py37"
-        if tf_full_version == TensorFlow._LATEST_1X_VERSION
-        else tests.integ.PYTHON_VERSION,
+        py_version=py_version,
         sagemaker_session=sagemaker_session,
         script_mode=True,
         # testing py-sdk functionality, no need to run against all TF versions
@@ -199,16 +193,14 @@ def test_mnist_async(sagemaker_session, cpu_instance_type, tf_full_version):
         _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name)
 
 
-def test_deploy_with_input_handlers(sagemaker_session, instance_type, tf_full_version):
+def test_deploy_with_input_handlers(sagemaker_session, instance_type, tf_full_version, py_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="py37"
-        if tf_full_version == TensorFlow._LATEST_1X_VERSION
-        else tests.integ.PYTHON_VERSION,
+        py_version=py_version,
         sagemaker_session=sagemaker_session,
         script_mode=True,
         framework_version=tf_full_version,