diff --git a/src/sagemaker/tensorflow/model.py b/src/sagemaker/tensorflow/model.py index c06fe74887..f2e82507e8 100644 --- a/src/sagemaker/tensorflow/model.py +++ b/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 ) diff --git a/tests/conftest.py b/tests/conftest.py index ceb2a03f51..db890d1a14 100644 --- a/tests/conftest.py +++ b/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) ) diff --git a/tests/data/tensorflow_mnist/mnist_v2.py b/tests/data/tensorflow_mnist/mnist_v2.py index 05467dee49..9efb282c49 100644 --- a/tests/data/tensorflow_mnist/mnist_v2.py +++ b/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__": diff --git a/tests/integ/sagemaker/workflow/test_model_create_and_registration.py b/tests/integ/sagemaker/workflow/test_model_create_and_registration.py index 0bdbc18c99..fd981b74e4 100644 --- a/tests/integ/sagemaker/workflow/test_model_create_and_registration.py +++ b/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( diff --git a/tests/integ/sagemaker/workflow/test_model_steps.py b/tests/integ/sagemaker/workflow/test_model_steps.py index da63bca597..089cdaf08f 100644 --- a/tests/integ/sagemaker/workflow/test_model_steps.py +++ b/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( diff --git a/tests/integ/sagemaker/workflow/test_training_steps.py b/tests/integ/sagemaker/workflow/test_training_steps.py index 181167ab31..bcff221afe 100644 --- a/tests/integ/sagemaker/workflow/test_training_steps.py +++ b/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" diff --git a/tests/integ/test_transformer.py b/tests/integ/test_transformer.py index d25f45d4db..8c99854d14 100644 --- a/tests/integ/test_transformer.py +++ b/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(