TF-2.16 test modification and handling (#4830)

* [DO NOT MERGE] Experimenting TF image_uri configs

* Add logic to Tf estimator

* Comment previous data

* Test with making changes to model.py

* Add override FW version

* Make changes to prevent non 2.16 from making the change

* Add Print

* Use hasattr

* Fix print and use getattr

* Print statement

* Only define override when needed

* Change to net.export() for TF2.16

* Skip tests failing due to TF-IO

* Reformatting

* Reformatting pylint

* Change tf_full_vesion fixture

* Add Version

* Revert unit test changes

* Preventing re-initialization of Version

* Revert config JSON changes

* Handle in case inf and training have different major.minor

* Introduce return version concept in tests

* Add TF2.16.1 inf to config

* Revert temp changes

---------

Co-authored-by: Erick Benitez-Ramos <[email protected]>

Author: shantanutrip

src/sagemaker/tensorflow/model.py

@@ -199,7 +199,7 @@ def __init__(
         # patch versions, but end up hosting the model of same TF version. For eg., the upstream
         # TFS-2.12.0 release was a bad release and hence a new TFS-2.12.1 release was made to host
         # models from TF-2.12.0.
-        training_inference_version_mismatch_dict = {"2.12.0": "2.12.1"}
+        training_inference_version_mismatch_dict = {"2.12.0": "2.12.1", "2.16.2": "2.16.1"}
         self.inference_framework_version = training_inference_version_mismatch_dict.get(
             framework_version, framework_version
         )

tests/conftest.py

@@ -22,6 +22,7 @@
 
 from botocore.config import Config
 from packaging.version import Version
+from packaging.specifiers import SpecifierSet
 
 from sagemaker import Session, image_uris, utils, get_execution_role
 from sagemaker.local import LocalSession
@@ -555,11 +556,18 @@ def tf_full_version(tensorflow_training_latest_version, tensorflow_inference_lat
     Fixture exists as such, since TF training and TFS have different latest versions.
     Otherwise, this would simply be a single latest version.
     """
-    return str(
-        min(
-            Version(tensorflow_training_latest_version),
-            Version(tensorflow_inference_latest_version),
-        )
+    tensorflow_training_latest_version = Version(tensorflow_training_latest_version)
+    tensorflow_inference_latest_version = Version(tensorflow_inference_latest_version)
+
+    return_version = min(
+        tensorflow_training_latest_version,
+        tensorflow_inference_latest_version,
+    )
+
+    return (
+        f"{return_version.major}.{return_version.minor}"
+        if return_version in SpecifierSet(">=2.16")
+        else str(return_version)
     )
 
 

tests/data/tensorflow_mnist/mnist_v2.py

@@ -198,7 +198,10 @@ def main(args):
 
         if args.current_host == args.hosts[0]:
             ckpt_manager.save()
-            net.save("/opt/ml/model/1")
+            if int(tf_major) > 2 or (int(tf_major) == 2 and int(tf_minor) >= 16):
+                net.export("/opt/ml/model/1")
+            else:
+                net.save("/opt/ml/model/1")
 
 
 if __name__ == "__main__":

tests/integ/sagemaker/workflow/test_model_create_and_registration.py

@@ -25,6 +25,9 @@
 
 import pytest
 
+from packaging.version import Version
+from packaging.specifiers import SpecifierSet
+
 from sagemaker.model_card.model_card import ModelCard, ModelOverview, ModelPackageModelCard
 from sagemaker.model_card.schema_constraints import ModelCardStatusEnum
 import tests
@@ -1250,6 +1253,11 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
     pipeline_name,
     region_name,
 ):
+    if Version(tf_full_version) in SpecifierSet("==2.16.*"):
+        pytest.skip(
+            "This test is failing in TensorFlow 2.16 beacuse of an upstream bug: "
+            "https://github.com/tensorflow/io/issues/2039"
+        )
     base_dir = os.path.join(DATA_DIR, "tensorflow_mnist")
     entry_point = os.path.join(base_dir, "mnist_v2.py")
     input_path = sagemaker_session_for_pipeline.upload_data(

tests/integ/sagemaker/workflow/test_model_steps.py

@@ -17,6 +17,9 @@
 
 import pytest
 
+from packaging.version import Version
+from packaging.specifiers import SpecifierSet
+
 from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker.workflow.fail_step import FailStep
 from sagemaker.workflow.functions import Join
@@ -589,6 +592,11 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics(
 def test_model_registration_with_tensorflow_model_with_pipeline_model(
     pipeline_session, role, tf_full_version, tf_full_py_version, pipeline_name
 ):
+    if Version(tf_full_version) in SpecifierSet("==2.16.*"):
+        pytest.skip(
+            "This test is failing in TensorFlow 2.16 beacuse of an upstream bug: "
+            "https://github.com/tensorflow/io/issues/2039"
+        )
     base_dir = os.path.join(DATA_DIR, "tensorflow_mnist")
     entry_point = os.path.join(base_dir, "mnist_v2.py")
     input_path = pipeline_session.upload_data(

tests/integ/sagemaker/workflow/test_training_steps.py

@@ -18,6 +18,9 @@
 
 import pytest
 
+from packaging.version import Version
+from packaging.specifiers import SpecifierSet
+
 from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker import TrainingInput, get_execution_role, utils, image_uris
 from sagemaker.debugger import (
@@ -235,6 +238,12 @@ def test_training_step_with_output_path_as_join(
 def test_tensorflow_training_step_with_parameterized_code_input(
     pipeline_session, role, tf_full_version, tf_full_py_version, pipeline_name
 ):
+    if Version(tf_full_version) in SpecifierSet("==2.16.*"):
+        pytest.skip(
+            "This test is failing in TensorFlow 2.16 beacuse of an upstream bug: "
+            "https://github.com/tensorflow/io/issues/2039"
+        )
+
     base_dir = os.path.join(DATA_DIR, "tensorflow_mnist")
     entry_point1 = "mnist_v2.py"
     entry_point2 = "mnist_dummy.py"

tests/integ/test_transformer.py

@@ -18,6 +18,9 @@
 
 import pytest
 
+from packaging.version import Version
+from packaging.specifiers import SpecifierSet
+
 from sagemaker import KMeans, s3, get_execution_role
 from sagemaker.mxnet import MXNet
 from sagemaker.pytorch import PyTorchModel
@@ -553,6 +556,12 @@ def test_transform_mxnet_logs(
 def test_transform_tf_kms_network_isolation(
     sagemaker_session, cpu_instance_type, tmpdir, tf_full_version, tf_full_py_version
 ):
+    if Version(tf_full_version) in SpecifierSet("==2.16.*"):
+        pytest.skip(
+            "This test is failing in TensorFlow 2.16 beacuse of an upstream bug: "
+            "https://github.com/tensorflow/io/issues/2039"
+        )
+
     data_path = os.path.join(DATA_DIR, "tensorflow_mnist")
 
     tf = TensorFlow(