Update xgboost dependency, execute in CI

.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": 4
+    "modification": 5
   }

.github/trigger_files/beam_PreCommit_Python_ML.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "revision": 22
+  "revision": 23
 }

CHANGES.md

@@ -95,10 +95,13 @@
 * (Python) `Regex.find_all` with `group=Regex.ALL` now returns every group, as `(match, group1, group2, ...)`, instead of only group 1, and with `outputEmpty=False` it drops a match only when the whole match is empty, matching the `group=0` behavior ([#40399](https://github.com/apache/beam/pull/40399)).
 * (Python) `Regex.all_matches` now returns every group of the regex, instead of failing with a `TypeError` when the regex has no groups and dropping nested or unmatched trailing groups ([#40562](https://github.com/apache/beam/issues/40562)).
 * (Python) `--dataflow_endpoint` now defaults to `''` instead of `https://dataflow.googleapis.com`. When unset, the Google Cloud client library resolves the Dataflow endpoint and honors `GOOGLE_CLOUD_UNIVERSE_DOMAIN`, so Dataflow jobs can be submitted in non-default universes. Code that reads `GoogleCloudOptions.dataflow_endpoint` should handle an empty value ([#40569](https://github.com/apache/beam/issues/40569)).
+* (Python) The `xgboost` extras now install `xgboost>=2.0.0,<4` and `datatable>=1.1.0,<2`, moving both the lower and upper bounds for both. The previous minimum `xgboost` and `datatable` versions (`1.6.0` and `1.0.0` respectively) are no longer allowed ([#40694](https://github.com/apache/beam/issues/40694)).
 
 ## Deprecations
 
 * Support for Python 3.10 was removed following its end-of-life ([#40250](https://github.com/apache/beam/issues/40250)).
+* (Python) `XGBoostModelHandlerDatatable` is deprecated and will be removed in Beam 2.79.0, along with the `datatable` dependency of the `xgboost` extra. `datatable` has not had a release since 2023 and has no wheels for Python 3.13+. Use `XGBoostModelHandlerPandas` (or the NumPy/SciPy handlers) instead ([#40705](https://github.com/apache/beam/issues/40705)).
+* X behavior is deprecated and will be removed in X versions ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Bugfixes
 

sdks/python/apache_beam/ml/inference/xgboost_inference.py

@@ -36,6 +36,7 @@
 from apache_beam.ml.inference.base import ModelT
 from apache_beam.ml.inference.base import PredictionResult
 from apache_beam.ml.inference.base import PredictionT
+from apache_beam.utils.annotations import deprecated
 
 __all__ = [
     'XGBoostModelHandler',
@@ -119,7 +120,7 @@ def __init__(
         before loading the model.
 
     **Supported Versions:** RunInference APIs in Apache Beam have been tested
-    with XGBoost 1.6.0 and 1.7.0
+    with XGBoost 2.x and 3.x
 
     XGBoost 1.0.0 introduced support for using JSON to save and load
     XGBoost models. XGBoost 1.6.0, additional support for Universal Binary JSON.
@@ -313,6 +314,12 @@ def get_num_bytes(self, batch: Sequence[scipy.sparse.csr_matrix]) -> int:
     return sum(sys.getsizeof(element) for element in batch)
 
 
+@deprecated(
+    since='2.78.0',
+    extra_message=(
+        'datatable has not had a release since 2023 and has no wheels for '
+        'Python 3.13+. XGBoostModelHandlerDatatable and the datatable '
+        'dependency will be removed in Beam 2.79.0.'))
 class XGBoostModelHandlerDatatable(XGBoostModelHandler[datatable.Frame,
                                                        PredictionResult,
                                                        Union[xgboost.Booster,

sdks/python/apache_beam/ml/inference/xgboost_inference_test.py

@@ -69,7 +69,9 @@ def _compare_keyed_prediction_result(
 
 
 def predict_fn(self, data):
-  self.inference_calls += 1
+  # Models loaded from disk inside the pipeline do not carry the counter set by
+  # build_monkeypatched_xgboost_classifier().
+  self.inference_calls = getattr(self, 'inference_calls', 0) + 1
   if isinstance(data, pandas.DataFrame):
     data = data.to_numpy()
   if isinstance(data, datatable.Frame):

sdks/python/apache_beam/ml/inference/xgboost_tests_requirements.txt

@@ -15,5 +15,5 @@
 # limitations under the License.
 #
 
-datatable==1.0.0
-xgboost>=1.6.0
+datatable>=1.1.0,<2
+xgboost>=2.0.0,<4

sdks/python/setup.py

@@ -166,6 +166,17 @@ def cythonize(*args, **kwargs):
 
 milvus_dependency = ['pymilvus>=2.5.10,<3.0.0']
 qdrant_dependency = ['qdrant-client>=1.15.0']
+# NOTE: When updating these bounds, also update the bounds in
+# apache_beam/ml/inference/xgboost_tests_requirements.txt, which is used by the
+# XGBoost RunInference integration tests, and the py311-xgboost-* environments
+# in tox.ini.
+xgboost_dependency = [
+    'xgboost>=2.0.0,<4',
+    # apache_beam.ml.inference.xgboost_inference imports datatable directly.
+    # 1.1.0 is the first release with Python 3.10+ wheels; no release has
+    # wheels for Python 3.13+.
+    'datatable>=1.1.0,<2',
+]
 
 # google-adk / OpenTelemetry require protobuf>=5; tensorflow-transform in
 # ml_test is pinned to versions that require protobuf<5 on Python 3.10. Those
@@ -608,13 +619,12 @@ def get_portability_package_data():
           'p310_ml_test': [
             'datatable',
           ] + ml_base + qdrant_dependency,
-          'p311_ml_test': [
-            'datatable',
-          ] + ml_base + qdrant_dependency,
+          # xgboost_dependency includes datatable, which only has wheels for
+          # Python <=3.12, so it is limited to the 3.11 and 3.12 test extras.
+          'p311_ml_test': ml_base + xgboost_dependency + qdrant_dependency,
           'p312_ml_test': [
-              'datatable',
               'typesafe-sdk>=0.7.1,<0.8',
-          ] + ml_base + qdrant_dependency,
+          ] + ml_base + xgboost_dependency + qdrant_dependency,
           # maintainer: milvus tests only run with this extension. Make sure it
           # is covered by docker-in-docker test when changing py version
           'p313_ml_test': ml_base_core + milvus_dependency + qdrant_dependency,
@@ -682,7 +692,7 @@ def get_portability_package_data():
               'transformers==4.25.1',  # tensorflow transient dep
               'absl-py>=0.12.0'
           ],
-          'xgboost': ['xgboost>=1.6.0,<2.1.3', 'datatable==1.0.0'],
+          'xgboost': xgboost_dependency,
           'tensorflow-hub': ['tensorflow-hub>=0.14.0,<0.16.0'],
           'milvus': milvus_dependency,
           'qdrant': qdrant_dependency,

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

@@ -158,6 +158,15 @@ toxTask "testPy311transformers-latest", "py311-transformers-latest", "${posargs}
 test.dependsOn "testPy311transformers-latest"
 postCommitPyDep.dependsOn "testPy311transformers-latest"
 
+// Create a test task for each major version of xgboost
+toxTask "testPy311xgboost-200", "py311-xgboost-200", "${posargs}"
+test.dependsOn "testPy311xgboost-200"
+postCommitPyDep.dependsOn "testPy311xgboost-200"
+
+toxTask "testPy311xgboost-300", "py311-xgboost-300", "${posargs}"
+test.dependsOn "testPy311xgboost-300"
+postCommitPyDep.dependsOn "testPy311xgboost-300"
+
 toxTask "testPy311embeddingsMLTransform", "py311-embeddings", "${posargs}"
 test.dependsOn "testPy311embeddingsMLTransform"
 postCommitPyDep.dependsOn "testPy311embeddingsMLTransform"

sdks/python/tox.ini

@@ -489,14 +489,12 @@ commands =
   # 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_tf {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
 
-[testenv:py311-xgboost-{160,170}]
+[testenv:py311-xgboost-{200,300}]
+# Keep these bounds in line with xgboost_dependency in setup.py.
 deps =
-  160:
-    xgboost>=1.6.0,<1.7.0
-    datatable==1.0.0
-  170:
-    xgboost>=1.7.0
-    datatable==1.0.0
+  200: xgboost>=2.0.0,<3.0.0
+  300: xgboost>=3.0.0,<4
+  datatable>=1.1.0,<2
 extras = test,gcp
 commands =
   # Log XGBoost version for debugging