Update pytorch dependencies

.github/trigger_files/beam_PostCommit_Python_Dependency.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run",
-    "modification": 6
+    "modification": 7
   }

sdks/python/setup.py

@@ -208,8 +208,8 @@ def cythonize(*args, **kwargs):
     # Drop this cap once transformers updates the CLIP call site to
     # `cls_token=` or tokenizers reinstates `cls=` as a deprecation alias.
     'tokenizers<0.23',
-    # Avoid torch 2.12.0+ which fails to run unit tests with segfault.
-    'torch<2.12.0',
+    # Keep in line with the py*-pytorch-* environments in tox.ini.
+    'torch>=2.3.0,<2.15.0',
     # Match tested transformers range.
     'transformers>=4.28.0,<4.56.0',
     # Keep tokenizers compatible with this transformers range.
@@ -659,16 +659,15 @@ def get_portability_package_data():
           # in https://github.com/apache/beam/blob/master/sdks/python/tox.ini
           # For more info, see
           # https://docs.google.com/document/d/1c84Gc-cZRCfrU8f7kWGsNR2o8oSRjCM-dGHO9KvPWPw/edit?usp=sharing
-          'torch': ['torch>=1.9.0,<2.8.0'],
+          'torch': ['torch>=2.3.0,<2.15.0'],
           'tensorflow': [
               'tensorflow>=2.12rc1,<2.22',  # tensorflow transitive dep
               'absl-py>=0.12.0'
           ],
           'transformers': [
               'transformers>=4.28.0,<4.56.0',
               'tensorflow>=2.12.0',
-              # Avoid torch 2.12.0+ which fails to run unit tests with segfault
-              'torch>=1.9.0,<2.12.0'
+              'torch>=2.3.0,<2.12.0'
           ],
           'ml_cpu': [
               'tensorflow>=2.12.0',

sdks/python/test-suites/tox/py311/build.gradle

@@ -96,24 +96,15 @@ postCommitPyDep.dependsOn "testPy311pandas-23"
 
 // TODO(https://github.com/apache/beam/issues/30908): Revise what are we testing
 
-// Create a test task for each minor version of pytorch
-
-toxTask "testPy311pytorch-111", "py311-pytorch-111", "${posargs}"
-test.dependsOn "testPy311pytorch-111"
-postCommitPyDep.dependsOn "testPy311pytorch-111"
-
-toxTask "testPy311pytorch-112", "py311-pytorch-112", "${posargs}"
-test.dependsOn "testPy311pytorch-112"
-postCommitPyDep.dependsOn "testPy311pytorch-112"
-
-toxTask "testPy311pytorch-113", "py311-pytorch-113", "${posargs}"
-test.dependsOn "testPy311pytorch-113"
-postCommitPyDep.dependsOn "testPy311pytorch-113"
-
-// run on precommit
-toxTask "testPy311pytorch-200", "py311-pytorch-200", "${posargs}"
-test.dependsOn "testPy311pytorch-200"
-postCommitPyDep.dependsOn "testPy311pytorch-200"
+// Create a test task for the lower and upper bounds of supported pytorch.
+// Keep in sync with [testenv:py311-pytorch-...] in tox.ini.
+toxTask "testPy311pytorch-23", "py311-pytorch-23", "${posargs}"
+test.dependsOn "testPy311pytorch-23"
+postCommitPyDep.dependsOn "testPy311pytorch-23"
+
+toxTask "testPy311pytorch-214", "py311-pytorch-214", "${posargs}"
+test.dependsOn "testPy311pytorch-214"
+postCommitPyDep.dependsOn "testPy311pytorch-214"
 
 toxTask "testPy311tft-113", "py311-tft-113", "${posargs}"
 test.dependsOn "testPy311tft-113"

sdks/python/test-suites/tox/py312/build.gradle

@@ -30,10 +30,6 @@ pythonVersion = '3.12'
 
 project.tasks.register("postCommitPyDep") {}
 
-toxTask "testPy312pytorch-271", "py312-pytorch-271", "${posargs}"
-test.dependsOn "testPy312pytorch-271"
-postCommitPyDep.dependsOn "testPy312pytorch-271"
-
 toxTask "testPy312tensorflow-216", "py312-tensorflow-216", "${posargs}"
 test.dependsOn "testPy312tensorflow-216"
 postCommitPyDep.dependsOn "testPy312tensorflow-216"

sdks/python/test-suites/tox/py314/build.gradle

@@ -30,3 +30,11 @@ project.tasks.register("postCommitPyDep") {}
 
 apply from: "../common.gradle"
 
+def posargs = project.findProperty("posargs") ?: ""
+
+// Latest supported pytorch. torch>=2.10 is required for Python 3.14.
+// Keep in sync with [testenv:py314-pytorch-214] in tox.ini.
+toxTask "testPy314pytorch-214", "py314-pytorch-214", "${posargs}"
+test.dependsOn "testPy314pytorch-214"
+postCommitPyDep.dependsOn "testPy314pytorch-214"
+

sdks/python/tox.ini

@@ -387,53 +387,34 @@ deps =
 commands =
   bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py'
 
-[testenv:py311-pytorch-{19,110,111,112,113}]
+[testenv:py311-pytorch-{23,214}]
+# Keep in line with the torch bounds in setup.py. torch<2.3 is not compatible
+# with numpy 2, so numpy is intentionally left unpinned.
 deps =
-  19: torch>=1.9.0,<1.10.0
-  110: torch>=1.10.0,<1.11.0
-  111: torch>=1.11.0,<1.12.0
-  112: torch>=1.12.0,<1.13.0
-  113: torch>=1.13.0,<1.14.0
-  numpy==1.26.4
-extras = test,gcp
-# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
-setenv =
-commands =
-  # Log torch version for debugging
-  /bin/sh -c "pip freeze | grep -E torch"
-  # Run all PyTorch<2 unit tests
-  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
-  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pytorch {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
-
-[testenv:py311-pytorch-200]
-deps =
-  200:
-    torch>=2.0.0,<2.1.0
-    mpmath==1.3.0
-    numpy==1.26.4
+  23: torch>=2.3.0,<2.4.0
+  214: torch>=2.14.0,<2.15.0
 extras = test,gcp
 # Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
 setenv =
 commands =
   # Log torch version for debugging
   /bin/sh -c "pip freeze | grep -E torch"
-  # Run all PyTorch>=2  unit tests
+  # Run all PyTorch unit tests
   # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
   /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pytorch {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
 
-[testenv:py{311,312}-pytorch-271]
+[testenv:py314-pytorch-214]
+# Keep in line with py311-pytorch-latest above. torch>=2.10 is required for
+# Python 3.14, so the torch 2.3 lower-bound env only runs on Python 3.11.
 deps =
-  200:
-    torch>=2.7.1,<2.8.0
-    mpmath==1.3.0
-    numpy==1.26.4
+  torch>=2.14.0,<2.15.0
 extras = test,gcp
 # Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
 setenv =
 commands =
   # Log torch version for debugging
   /bin/sh -c "pip freeze | grep -E torch"
-  # Run all PyTorch>=2  unit tests
+  # Run all PyTorch unit tests
   # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
   /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pytorch {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'