Adding python 3.11 support and Removing TF-Lite Dependency

.github/workflows/ci-test.yml

@@ -22,7 +22,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ['3.9', '3.10']
+        python-version: ['3.9', '3.10','3.11']
         which-tests: ["not e2e", "e2e"]
         dependency-selector: ["NIGHTLY", "DEFAULT"]
 

nightly_test_constraints.txt

@@ -125,7 +125,7 @@ greenlet==3.1.1
 grpc-google-iam-v1==0.13.1
 grpc-interceptor==0.15.4
 grpcio==1.66.2
-grpcio-status==1.48.2
+# grpcio-status==1.48.2
 gunicorn==23.0.0
 h11==0.14.0
 h5py==3.12.1
@@ -168,9 +168,9 @@ jupyterlab_server==2.27.3
 jupyterlab_widgets==1.1.10
 tf-keras==2.16.0
 keras-tuner==1.4.7
-kfp==2.5.0
-kfp-pipeline-spec==0.2.2
-kfp-server-api==2.0.5
+# kfp==2.5.0
+# kfp-pipeline-spec==0.2.2
+# kfp-server-api==2.0.5
 kt-legacy==1.0.5
 kubernetes==26.1.0
 lazy-object-proxy==1.10.0
@@ -246,7 +246,7 @@ promise==2.3
 prompt_toolkit==3.0.48
 propcache==0.2.0
 proto-plus==1.24.0
-protobuf==3.20.3
+# protobuf==3.20.3
 psutil==6.0.0
 ptyprocess==0.7.0
 pyarrow-hotfix==0.6
@@ -315,7 +315,7 @@ tensorflow-decision-forests==1.9.2
 tensorflow-estimator==2.15.0
 tensorflow-hub==0.15.0
 tensorflow-io==0.24.0
-tensorflow-io-gcs-filesystem==0.24.0
+# tensorflow-io-gcs-filesystem==0.24.0
 tensorflow-metadata>=1.16.1
 # tensorflow-ranking==0.5.5
 tensorflow-serving-api==2.16.1
@@ -327,7 +327,7 @@ tensorstore==0.1.66
 termcolor==2.5.0
 terminado==0.18.1
 text-unidecode==1.3
-tflite-support==0.4.4
+# tflite-support==0.4.4
 tfx-bsl>=1.16.1
 threadpoolctl==3.5.0
 time-machine==2.16.0

pyproject.toml

@@ -22,6 +22,7 @@ classifiers = [
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3 :: Only",
     "Topic :: Scientific/Engineering",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
@@ -31,7 +32,7 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules"
 ]
 keywords = ["tensorflow", "tfx"]
-requires-python = ">=3.9,<3.11"
+requires-python = ">=3.9,<3.12"
 [project.urls]
 Homepage = "https://www.tensorflow.org/tfx"
 Repository = "https://github.com/tensorflow/tfx"

test_constraints.txt

@@ -125,7 +125,7 @@ greenlet==3.1.1
 grpc-google-iam-v1==0.13.1
 grpc-interceptor==0.15.4
 grpcio==1.66.2
-grpcio-status==1.48.2
+# grpcio-status==1.48.2
 gunicorn==23.0.0
 h11==0.14.0
 h5py==3.12.1
@@ -168,9 +168,9 @@ jupyterlab_server==2.27.3
 jupyterlab_widgets==1.1.10
 tf-keras==2.16.0
 keras-tuner==1.4.7
-kfp==2.5.0
-kfp-pipeline-spec==0.2.2
-kfp-server-api==2.0.5
+# kfp==2.5.0
+# kfp-pipeline-spec==0.2.2
+# kfp-server-api==2.0.5
 kt-legacy==1.0.5
 kubernetes==26.1.0
 lazy-object-proxy==1.10.0
@@ -246,7 +246,7 @@ promise==2.3
 prompt_toolkit==3.0.48
 propcache==0.2.0
 proto-plus==1.24.0
-protobuf==3.20.3
+# protobuf==3.20.3
 psutil==6.0.0
 ptyprocess==0.7.0
 pyarrow-hotfix==0.6
@@ -315,7 +315,7 @@ tensorflow-decision-forests==1.9.2
 tensorflow-estimator==2.15.0
 tensorflow-hub==0.15.0
 tensorflow-io==0.24.0
-tensorflow-io-gcs-filesystem==0.24.0
+# tensorflow-io-gcs-filesystem==0.24.0
 tensorflow-metadata>=1.16.1
 # tensorflow-ranking==0.5.5
 tensorflow-serving-api==2.16.1
@@ -327,7 +327,7 @@ tensorstore==0.1.66
 termcolor==2.5.0
 terminado==0.18.1
 text-unidecode==1.3
-tflite-support==0.4.4
+# tflite-support==0.4.4
 tfx-bsl>=1.16.1
 threadpoolctl==3.5.0
 time-machine==2.16.0

tfx/components/infra_validator/request_builder_test.py

@@ -440,7 +440,7 @@ def setUp(self):
   def _PrepareTFServingRequestBuilder(self):
     patcher = mock.patch.object(
         request_builder, '_TFServingRpcRequestBuilder',
-        wraps=request_builder._TFServingRpcRequestBuilder)
+        autospec=True)
     builder_cls = patcher.start()
     self.addCleanup(patcher.stop)
     return builder_cls
@@ -466,7 +466,7 @@ def testBuildRequests_TFServing(self):
         model_name='foo',
         signatures={'serving_default': mock.ANY})
     builder.ReadExamplesArtifact.assert_called_with(
-        self._examples,
+        examples=self._examples,
         split_name='eval',
         num_examples=1)
     builder.BuildRequests.assert_called()

tfx/components/trainer/rewriting/rewriter_factory.py

@@ -21,13 +21,12 @@
 
 from tfx.components.trainer.rewriting import rewriter
 
-TFLITE_REWRITER = 'TFLiteRewriter'
+# TFLITE_REWRITER = 'TFLiteRewriter'
 TFJS_REWRITER = 'TFJSRewriter'
 
 
-def _load_tflite_rewriter():
-  importlib.import_module('tfx.components.trainer.rewriting.tflite_rewriter')
-
+# def _load_tflite_rewriter():
+#  importlib.import_module('tfx.components.trainer.rewriting.tflite_rewriter')
 
 def _load_tfjs_rewriter():
   try:
@@ -43,7 +42,7 @@ def _load_tfjs_rewriter():
 class _RewriterFactory:
   """Factory class for rewriters."""
   _LOADERS = {
-      TFLITE_REWRITER.lower(): _load_tflite_rewriter,
+     # TFLITE_REWRITER.lower(): _load_tflite_rewriter,
       TFJS_REWRITER.lower(): _load_tfjs_rewriter,
   }
   _loaded = set()
@@ -55,14 +54,14 @@ def _maybe_load_public_rewriter(cls, lower_rewriter_type: str):
       cls._LOADERS[lower_rewriter_type]()
       cls._loaded.add(lower_rewriter_type)
 
-  @classmethod
-  def get_rewriter_cls(cls, rewriter_type: str):
-    rewriter_type = rewriter_type.lower()
-    cls._maybe_load_public_rewriter(rewriter_type)
-    for subcls in rewriter.BaseRewriter.__subclasses__():
-      if subcls.__name__.lower() == rewriter_type:
-        return subcls
-    raise ValueError('Failed to find rewriter: {}'.format(rewriter_type))
+#  @classmethod
+#  def get_rewriter_cls(cls, rewriter_type: str):
+#    rewriter_type = rewriter_type.lower()
+#    cls._maybe_load_public_rewriter(rewriter_type)
+#    for subcls in rewriter.BaseRewriter.__subclasses__():
+#      if subcls.__name__.lower() == rewriter_type:
+#        return subcls
+#    raise ValueError('Failed to find rewriter: {}'.format(rewriter_type))
 
 
 def create_rewriter(rewriter_type: str, *args,

tfx/components/trainer/rewriting/rewriter_factory_test.py

@@ -31,13 +31,13 @@ def _tfjs_installed():
 
 class RewriterFactoryTest(parameterized.TestCase):
 
-  @parameterized.named_parameters(
-      ('TFLite', rewriter_factory.TFLITE_REWRITER))
-  def testRewriterFactorySuccessfullyCreated(self, rewriter_name):
-    tfrw = rewriter_factory.create_rewriter(rewriter_name, name='my_rewriter')
-    self.assertTrue(tfrw)
-    self.assertEqual(type(tfrw).__name__, rewriter_name)
-    self.assertEqual(tfrw.name, 'my_rewriter')
+  # @parameterized.named_parameters(
+    #  ('TFLite', rewriter_factory.TFLITE_REWRITER))
+  # def testRewriterFactorySuccessfullyCreated(self, rewriter_name):
+  #  tfrw = rewriter_factory.create_rewriter(rewriter_name, name='my_rewriter')
+  #  self.assertTrue(tfrw)
+  #  self.assertEqual(type(tfrw).__name__, rewriter_name)
+  #  self.assertEqual(tfrw.name, 'my_rewriter')
 
   @unittest.skipUnless(_tfjs_installed(), 'tensorflowjs is not installed')
   def testRewriterFactorySuccessfullyCreatedTFJSRewriter(self):