Upgrade TFMA to TF 2.21.0 and Added Python 3.12/3.13 support

.github/workflows/ci-test.yml

@@ -16,7 +16,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ['3.9', '3.10', '3.11']
+        python-version: ['3.10', '3.11', '3.12', '3.13']
 
     steps:
     - name: Checkout repository

RELEASE.md

@@ -4,9 +4,38 @@
 
 ## Major Features and Improvements
 
+*   Adds official support for Python 3.12 and 3.13.
+*   Drops support for Python 3.9.
+*   Depends on `tensorflow>=2.21.0,<2.22.0`.
+*   Depends on `protobuf>=6.31.1,<7.0.0`.
+*   Depends on `pyarrow>14`.
+*   Updates the minimum Bazel version required to build TFMA to 7.4.1.
+
 ## Bug fixes and other Changes
 
+*   **Pickling Stability Architecture**:
+    *   Refactored core test modules (`rouge_test.py`, `stats_test.py`, `model_util_test.py`, `confidence_intervals_util_test.py`, `metrics_plots_and_validations_evaluator_test.py`) to use a new **class-based matcher architecture**.
+    *   Replaced nested closures with module-level classes (e.g., `CheckResult`, `CheckResultMean`) to ensure full serializability for `PrismRunner` on Python 3.13.
+    *   Removed `self` (test instance) capture in Beam matchers to resolve `RuntimeError: Unable to pickle fn` during distributed execution.
+    *   Enabled `--no_save_main_session` for all Beam pipelines in the test suite to prevent unintentional serialization of the main session and shared resources.
+*   **NumPy 2.0 & Python 3.13 Compatibility**:
+    *   Standardized on safe scalar extraction by replacing `float(ndarray)` with `.item()` or `metric_util.safe_to_scalar` in aggregation, attributions, calibration, flip metrics, and NDCG modules to resolve `TypeError` in Beam pipelines.
+    *   Fixed a batching bug in `flip_metrics.py` to correctly process all examples in a Beam batch.
+    *   Implemented robust, warning-free division in AUC, PR AUC, and confusion matrix calculations using `np.divide` with `where` clauses to prevent `RuntimeWarning`.
+    *   Fixed `TypeError` in `poisson_bootstrap.py` by removing redundant size argument from `poisson(1, 1)` to return a scalar.
+    *   Implemented `kind='stable'` sort in `metric_util.top_k_indices` for deterministic tie-breaking.
+*   **Bug Fixes and Functional Corrections**:
+    *   Fixed a critical regression in `metric_util.py` where `SubKey(k=k)` incorrectly selected the first prediction instead of the requested k-th largest prediction.
+    *   Fixed `UnparsedFlagAccessError` in `ModelSignaturesDoFn` tests by removing direct `absl.flags` access in pickling-sensitive contexts.
+    *   Removed obsolete `@unittest.expectedFailure` decorators from tests that are now passing in the stabilized environment.
+    *   Fixed various indentation and syntax errors in utility tests.
+    *   Improved virtual environment relocation strategy to resolve Bazel sandbox access issues for `numpy` and other C-extension headers.
+    *   Fixed `false_omission_rate` in `binary_confusion_matrices.py` to return NaN when undefined, resolving proto mismatches in `confusion_matrix_plot_test.py` and `score_distribution_plot_test.py`.
+    *   Fixed `NotFoundError` in `model_eval_lib_test.py` by ensuring temporary directories exist before writing files using `tf.io.gfile.makedirs`.
+    *   Added missing `numpy` imports in Beam-based modules to fix `NameError` regressions.
+
 ## Breaking Changes
+*   Python 3.9 is no longer supported. The minimum supported Python version is now 3.10.
 
 ## Deprecations
 

WORKSPACE

@@ -2,21 +2,22 @@ workspace(name = "org_tensorflow_model_analysis")
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-# TF 2.17.1
+# TF 2.21.0
 # LINT.IfChange(tf_commit)
-_TENSORFLOW_GIT_COMMIT = "3c92ac03cab816044f7b18a86eb86aa01a294d95"
+_TENSORFLOW_GIT_COMMIT = "a481b10260dfdf833a1b16007eead49c1d7febf3"
 
 # LINT.ThenChange(:io_bazel_rules_closure)
 http_archive(
     name = "org_tensorflow",
-    sha256 = "317dd95c4830a408b14f3e802698eb68d70d81c7c7cfcd3d28b0ba023fe84a68",
+    sha256 = "ef3568bb4865d6c1b2564fb5689c19b6b9a5311572cd1f2ff9198636a8520921",
     strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
     urls = [
-        "https://mirror.bazel.build/github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
+        "http://mirror.tensorflow.org/github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
         "https://github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
     ],
 )
 
+
 # Needed by tensorboard. Because these http_archives do not handle transitive
 # dependencies, we need to unroll them here.
 http_archive(
@@ -76,21 +77,22 @@ http_archive(
 
 load("@org_tensorflow_tensorboard//third_party:workspace.bzl", "tensorboard_workspace")
 
-_PROTOBUF_COMMIT = "4.25.6"  # 4.25.6
+_PROTOBUF_COMMIT = "74211c0dfc2777318ab53c2cd2c317a2ef9012de"
 
 http_archive(
     name = "com_google_protobuf",
-    sha256 = "ff6e9c3db65f985461d200c96c771328b6186ee0b10bc7cb2bbc87cf02ebd864",
+    sha256 = "554e847e46c705bfc44fb2d0ae5bf78f34395fcbfd86ba747338b570eef26771",
     strip_prefix = "protobuf-%s" % _PROTOBUF_COMMIT,
     urls = [
-        "https://github.com/protocolbuffers/protobuf/archive/v4.25.6.zip",
+        "https://github.com/protocolbuffers/protobuf/archive/%s.zip" % _PROTOBUF_COMMIT,
     ],
 )
 
 load("@com_google_protobuf//:protobuf_deps.bzl", "protobuf_deps")
 
 protobuf_deps()
 
+
 tensorboard_workspace()
 
 load("//third_party:workspace.bzl", "tensorflow_model_analysis_workspace")
@@ -100,4 +102,4 @@ tensorflow_model_analysis_workspace()
 
 load("@bazel_skylib//lib:versions.bzl", "versions")
 
-versions.check("6.5.0")
+versions.check("7.4.1")

setup.py

@@ -319,21 +319,19 @@ def select_constraint(default, nightly=None, git_master=None):
     "install_requires": [
         # Sort alphabetically
         "absl-py>=0.9,<2.0.0",
-        'apache-beam[gcp]>=2.53,<3;python_version>="3.11"',
-        'apache-beam[gcp]>=2.50,<2.51;python_version<"3.11"',
+        "apache-beam[gcp]>=2.53,<3",
         "ipython>=7,<8",
         "ipywidgets>=7,<8",
         "numpy>=1.23.5",
-        "pandas>=1.0,<2",
+        "pandas>=1.0,<3",
         "pillow>=9.4.0",
-        'protobuf>=4.25.2,<6.0.0;python_version>="3.11"',
-        'protobuf>=4.21.6,<6.0.0;python_version<"3.11"',
-        "pyarrow>=10,<11",
+        "protobuf>=6.31.1",
+        "pyarrow>14",
         "rouge-score>=0.1.2,<2",
         "sacrebleu>=2.3,<4",
         "scipy>=1.4.1,<2",
         "six>=1.12,<2",
-        "tensorflow>=2.17,<2.18",
+        "tensorflow>=2.21.0",
         "tensorflow-estimator>=2.10",
         "tensorflow-metadata"
         + select_constraint(
@@ -343,17 +341,17 @@ def select_constraint(default, nightly=None, git_master=None):
         ),
         "tfx-bsl"
         + select_constraint(
-            default=">=1.17.1,<1.18.0",
+            default="@git+https://github.com/vkarampudi/tfx-bsl@testing",
             nightly=">=1.18.0.dev",
-            git_master="@git+https://github.com/tensorflow/tfx-bsl@master",
+            git_master="@git+https://github.com/vkarampudi/tfx-bsl@testing",
         ),
         "tf-keras",
     ],
     "extras_require": {
         "all": [*_make_extra_packages_tfjs(), *_make_docs_packages()],
         "docs": _make_docs_packages(),
     },
-    "python_requires": ">=3.9,<4",
+    "python_requires": ">=3.10,<3.14",
     "packages": find_packages(),
     "zip_safe": False,
     "cmdclass": {
@@ -375,9 +373,10 @@ def select_constraint(default, nightly=None, git_master=None):
         "Operating System :: OS Independent",
         "Programming Language :: Python",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Programming Language :: Python :: 3.13",
         "Programming Language :: Python :: 3 :: Only",
         "Topic :: Scientific/Engineering",
         "Topic :: Scientific/Engineering :: Mathematics",

tensorflow_model_analysis/api/model_eval_lib_test.py

@@ -1383,34 +1383,31 @@ def testRunModelAnalysisWithSchema(self):
         self.assertEqual(1.0, got_buckets[1]["lowerThresholdInclusive"])
         self.assertEqual(2.0, got_buckets[-2]["upperThresholdExclusive"])
 
-    # PR 189: Remove the `expectedFailure` mark if the test passes
-    @unittest.expectedFailure
     def testLoadValidationResult(self):
         result = validation_result_pb2.ValidationResult(validation_ok=True)
         path = os.path.join(absltest.get_default_test_tmpdir(), "results.tfrecord")
+        tf.io.gfile.makedirs(os.path.dirname(path))
         with tf.io.TFRecordWriter(path) as writer:
             writer.write(result.SerializeToString())
         loaded_result = model_eval_lib.load_validation_result(path)
         self.assertTrue(loaded_result.validation_ok)
 
-    # PR 189: Remove the `expectedFailure` mark if the test passes
-    @unittest.expectedFailure
     def testLoadValidationResultDir(self):
         result = validation_result_pb2.ValidationResult(validation_ok=True)
         path = os.path.join(
             absltest.get_default_test_tmpdir(), constants.VALIDATIONS_KEY
         )
+        tf.io.gfile.makedirs(os.path.dirname(path))
         with tf.io.TFRecordWriter(path) as writer:
             writer.write(result.SerializeToString())
         loaded_result = model_eval_lib.load_validation_result(os.path.dirname(path))
         self.assertTrue(loaded_result.validation_ok)
 
-    # PR 189: Remove the `expectedFailure` mark if the test passes
-    @unittest.expectedFailure
     def testLoadValidationResultEmptyFile(self):
         path = os.path.join(
             absltest.get_default_test_tmpdir(), constants.VALIDATIONS_KEY
         )
+        tf.io.gfile.makedirs(os.path.dirname(path))
         with tf.io.TFRecordWriter(path):
             pass
         with self.assertRaises(AssertionError):
@@ -1538,15 +1535,16 @@ def testBytesProcessedCountForSerializedExamples(self):
         ]
         serialized_examples = [example.SerializeToString() for example in examples]
         expected_num_bytes = sum([len(se) for se in serialized_examples])
-        with beam.Pipeline() as p:
-            _ = (
-                p
-                | beam.Create(serialized_examples)
-                | "InputsToExtracts" >> model_eval_lib.InputsToExtracts()
-                | "ExtractAndEvaluate"
-                >> model_eval_lib.ExtractAndEvaluate(extractors=[], evaluators=[])
-            )
+        p = beam.Pipeline()
+        _ = (
+            p
+            | beam.Create(serialized_examples)
+            | "InputsToExtracts" >> model_eval_lib.InputsToExtracts()
+            | "ExtractAndEvaluate"
+            >> model_eval_lib.ExtractAndEvaluate(extractors=[], evaluators=[])
+        )
         pipeline_result = p.run()
+        pipeline_result.wait_until_finish()
         metrics = pipeline_result.metrics()
         actual_counter = metrics.query(
             beam.metrics.metric.MetricsFilter().with_name("extract_input_bytes")
@@ -1566,15 +1564,16 @@ def testBytesProcessedCountForRecordBatches(self):
         decoder = example_coder.ExamplesToRecordBatchDecoder()
         record_batch = decoder.DecodeBatch(examples)
         expected_num_bytes = record_batch.nbytes
-        with beam.Pipeline() as p:
-            _ = (
-                p
-                | beam.Create(record_batch)
-                | "BatchedInputsToExtracts" >> model_eval_lib.BatchedInputsToExtracts()
-                | "ExtractAndEvaluate"
-                >> model_eval_lib.ExtractAndEvaluate(extractors=[], evaluators=[])
-            )
+        p = beam.Pipeline()
+        _ = (
+            p
+            | beam.Create(record_batch)
+            | "BatchedInputsToExtracts" >> model_eval_lib.BatchedInputsToExtracts()
+            | "ExtractAndEvaluate"
+            >> model_eval_lib.ExtractAndEvaluate(extractors=[], evaluators=[])
+        )
         pipeline_result = p.run()
+        pipeline_result.wait_until_finish()
         metrics = pipeline_result.metrics()
         actual_counter = metrics.query(
             beam.metrics.metric.MetricsFilter().with_name("extract_input_bytes")