use tensorflow-transform 1.16

Author: aIbrahiim

sdks/python/setup.py

@@ -541,7 +541,7 @@ def get_portability_package_data():
               # tensorflow-transform requires dill, but doesn't set dill as a
               # hard requirement in setup.py.
               'dill',
-              'tensorflow-transform',
+              'tensorflow-transform>=1.16.0,<1.17.0',
               # Comment out xgboost as it is breaking presubmit python ml
               # tests due to tag check introduced since pip 24.2
               # https://github.com/apache/beam/issues/31285

sdks/python/tox.ini

@@ -105,12 +105,21 @@ commands =
 
 [testenv:py{310,311}-ml]
 # Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
+# Pin versions to avoid pip backtracking (2+ hours) on self-hosted runners.
+# tensorflow 2.17 needs protobuf<5; grpcio-status 1.65 needs protobuf>=5.
+# tensorflow-transform 1.17 requires apache-beam<2.51 for py<3.11. Use tft 1.16
+# which supports beam>=2.47. grpcio 1.62.2 + protobuf 4.x for compatibility.
 deps =
   pip==25.0.1
   accelerate>=1.6.0
-  onnx<1.19.0
+  protobuf>=4.25,<5
+  grpcio==1.62.2
+  grpcio-status==1.62.2
+  tensorflow>=2.16,<2.17
+pip_pre = False
 setenv =
 extras = test,gcp,dataframe,ml_test
+install_command = {envbindir}/python -m pip install --retries 10 --extra-index-url https://download.pytorch.org/whl/cpu {opts} {packages}
 commands =
   # Log tensorflow version for debugging
   /bin/sh -c "pip freeze | grep -E tensorflow"