TF-2.16 test modification and handling

src/sagemaker/image_uri_config/tensorflow.json

@@ -331,7 +331,8 @@
             "2.11": "2.11.1",
             "2.12": "2.12.1",
             "2.13": "2.13.0",
-            "2.14": "2.14.1"
+            "2.14": "2.14.1",
+            "2.16": "2.16.1"
         },
         "versions": {
             "1.4.1": {
@@ -2188,6 +2189,48 @@
                     "us-west-2": "763104351884"
                 },
                 "repository": "tensorflow-inference"
+            },
+            "2.16.1": {
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-south-2": "772153158452",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ap-southeast-4": "457447274322",
+                    "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
+                    "us-isof-south-1": "454834333376",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "repository": "tensorflow-inference"
             }
         }
     },

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,10 +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.
     """
+    tensorflow_training_latest_version = Version(tensorflow_training_latest_version)
+    tensorflow_inference_latest_version = Version(tensorflow_inference_latest_version)
+
+    if tensorflow_training_latest_version in SpecifierSet(">=2.16"):
+        return (
+            f"{tensorflow_training_latest_version.major}"
+            f".{tensorflow_training_latest_version.minor}"
+        )
     return str(
         min(
-            Version(tensorflow_training_latest_version),
-            Version(tensorflow_inference_latest_version),
+            tensorflow_training_latest_version,
+            tensorflow_inference_latest_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(