Revert unit test changes

Author: shantanutrip

src/sagemaker/estimator.py

@@ -3727,12 +3727,11 @@ def training_image_uri(self, region=None):
         Returns:
             str: The URI of the Docker image.
         """
+
         return image_uris.get_training_image_uri(
             region=region or self.sagemaker_session.boto_region_name,
             framework=self._framework_name,
-            framework_version=getattr(
-                self, "override_fw_version", self.framework_version  # pylint: disable=no-member
-            ),
+            framework_version=self.framework_version,  # pylint: disable=no-member
             py_version=self.py_version,  # pylint: disable=no-member
             image_uri=self.image_uri,
             distribution=getattr(self, "distribution", None),

src/sagemaker/tensorflow/estimator.py

@@ -177,11 +177,6 @@ def __init__(
                 fw.python_deprecation_warning(self._framework_name, defaults.LATEST_PY2_VERSION)
             )
         self.framework_version = framework_version
-        # if self.framework_version and "2.16" in self.framework_version:
-        #     self.override_fw_version = "2.16"
-        #     # TODO: Revert
-        #     print(f"ABCD123 setting self.override_fw_version to {self.override_fw_version}")
-
         self.py_version = py_version
         self.instance_type = instance_type
 

tests/conftest.py

@@ -557,7 +557,10 @@ def tf_full_version(tensorflow_training_latest_version, tensorflow_inference_lat
     Otherwise, this would simply be a single latest version.
     """
     if Version(tensorflow_training_latest_version) in SpecifierSet(">=2.16"):
-        return f"{Version(tensorflow_training_latest_version).major}.{Version(tensorflow_training_latest_version).minor}"
+        return (
+            f"{Version(tensorflow_training_latest_version).major}"
+            f".{Version(tensorflow_training_latest_version).minor}"
+        )
     return str(
         min(
             Version(tensorflow_training_latest_version),

tests/unit/sagemaker/tensorflow/test_estimator.py

@@ -18,8 +18,6 @@
 
 from mock import patch, Mock, MagicMock
 from packaging import version
-from packaging.version import Version
-from packaging.specifiers import SpecifierSet
 import pytest
 
 from sagemaker.estimator import _TrainingJob
@@ -556,14 +554,9 @@ def test_fit_mwms(
 
     expected_train_args = _create_train_job(framework_version, py_version=py_version)
     expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs
-    if Version(framework_version) in SpecifierSet("==2.16.*"):
-        expected_train_args["image_uri"] = (
-            f"763104351884.dkr.ecr.{REGION}.amazonaws.com/tensorflow-training:2.16-cpu-{py_version}"
-        )
-    else:
-        expected_train_args["image_uri"] = (
-            f"763104351884.dkr.ecr.{REGION}.amazonaws.com/tensorflow-training:{framework_version}-cpu-{py_version}"
-        )
+    expected_train_args["image_uri"] = (
+        f"763104351884.dkr.ecr.{REGION}.amazonaws.com/tensorflow-training:{framework_version}-cpu-{py_version}"
+    )
     expected_train_args["job_name"] = f"tensorflow-training-{TIMESTAMP}"
     expected_train_args["hyperparameters"][TensorFlow.LAUNCH_MWMS_ENV_NAME] = json.dumps(True)
     expected_train_args["hyperparameters"]["sagemaker_job_name"] = json.dumps(

tests/unit/test_processing.py

@@ -18,8 +18,6 @@
 import pytest
 from mock import Mock, patch, MagicMock
 from packaging import version
-from packaging.version import Version
-from packaging.specifiers import SpecifierSet
 
 from sagemaker import LocalSession
 from sagemaker.dataset_definition.inputs import (
@@ -509,14 +507,7 @@ def test_tensorflow_processor_with_required_parameters(
     else:
         tensorflow_image_uri = (
             "763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:{}-cpu-{}"
-        ).format(
-            (
-                "2.16"
-                if Version(tensorflow_training_version) in SpecifierSet("==2.16.*")
-                else tensorflow_training_version
-            ),
-            tensorflow_training_py_version,
-        )
+        ).format(tensorflow_training_version, tensorflow_training_py_version)
 
     expected_args["app_specification"]["ImageUri"] = tensorflow_image_uri