core/validate.py#L176 checks the current TensorFlow Version on your local machine. It then asserts that the value is in the list from get_cloud_tpu_supported_tf_version which contains the list [“2.1”]. This makes it currently impossible to perform TPU training (from what I can tell, you can’t get a TF version where tf.__version__ is exactly “2.1” with no postpended subversion.

Removes this check so TPU training will work again

PiperOrigin-RevId: 388815387

Author: LukeWood

src/python/tensorflow_cloud/core/tests/unit/validate_test.py

@@ -16,7 +16,6 @@
 import os
 
 from absl.testing import absltest
-import mock
 
 from tensorflow_cloud.core import machine_config
 from tensorflow_cloud.core import validate
@@ -302,30 +301,5 @@ def test_invalid_tpu_accelerator_count(self):
                 called_from_notebook=False,
             )
 
-    @mock.patch("tensorflow_cloud.utils.tf_utils.get_version")  # pylint: disable=line-too-long
-    def test_invalid_tpu_accelerator_tf_version(self, mock_get_version):
-        mock_get_version.return_value = "2.2.0"
-        with self.assertRaisesRegex(
-            NotImplementedError,
-            r"TPUs are only supported for TF version <= 2.1.0",
-        ):
-            validate.validate(
-                entry_point=None,
-                distribution_strategy="auto",
-                requirements_txt=None,
-                chief_config=machine_config.COMMON_MACHINE_CONFIGS["CPU"],
-                worker_config=machine_config.MachineConfig(
-                    accelerator_type=machine_config.AcceleratorType.TPU_V2,
-                    accelerator_count=8,
-                ),
-                worker_count=1,
-                entry_point_args=None,
-                stream_logs=True,
-                docker_image_build_bucket=None,
-                called_from_notebook=False,
-                docker_parent_image=None,
-            )
-
-
 if __name__ == "__main__":
     absltest.main()

src/python/tensorflow_cloud/core/validate.py

@@ -17,7 +17,6 @@
 
 from . import gcp
 from . import machine_config
-from ..utils import tf_utils
 
 
 def validate(
@@ -169,16 +168,6 @@ def _validate_cluster_config(
                 "Expected worker_count=1 for TPU `worker_config`. "
                 "Received {}.".format(worker_count)
             )
-        elif docker_parent_image is None:
-            # If the user has not provided a custom Docker image, then verify
-            # that the TF version is compatible with Cloud TPU support.
-            # https://cloud.google.com/ai-platform/training/docs/runtime-version-list#tpu-support  # pylint: disable=line-too-long
-            version = tf_utils.get_version()
-            if (version is not None and
-                version not in gcp.get_cloud_tpu_supported_tf_versions()):
-                raise NotImplementedError(
-                    "TPUs are only supported for TF version <= 2.1.0"
-                )
 
 
 def _validate_job_labels(job_labels):