Merge branch 'smp-2.5' of https://github.com/adtian2/sagemaker-python-sdk into smp-2.5

Author: adtian2

src/sagemaker/image_uri_config/huggingface-llm-neuronx.json

@@ -4,7 +4,7 @@
             "inf2"
         ],
         "version_aliases": {
-            "0.0": "0.0.23"
+            "0.0": "0.0.24"
         },
         "versions": {
             "0.0.16": {
@@ -238,6 +238,35 @@
                 "container_version": {
                     "inf2": "ubuntu22.04"
                 }
+            },
+            "0.0.24": {
+                "py_versions": [
+                    "py310"
+                ],
+                "registries": {
+                    "ap-northeast-1": "763104351884",
+                    "ap-south-1": "763104351884",
+                    "ap-south-2": "772153158452",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-4": "457447274322",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572"
+                },
+                "tag_prefix": "2.1.2-optimum0.0.24",
+                "repository": "huggingface-pytorch-tgi-inference",
+                "container_version": {
+                    "inf2": "ubuntu22.04"
+                }
             }
         }
     }

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(

tests/unit/sagemaker/image_uris/test_huggingface_llm.py

@@ -55,6 +55,7 @@
         "0.0.21": "1.13.1-optimum0.0.21-neuronx-py310-ubuntu22.04",
         "0.0.22": "2.1.2-optimum0.0.22-neuronx-py310-ubuntu22.04",
         "0.0.23": "2.1.2-optimum0.0.23-neuronx-py310-ubuntu22.04",
+        "0.0.24": "2.1.2-optimum0.0.24-neuronx-py310-ubuntu22.04",
     },
 }