add barrier before creating quantile dmatrix

wbo4958 · May 30, 2024 · b01c8e2 · b01c8e2
1 parent e414c2c
commit b01c8e2
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python-package/xgboost/spark/data.py b/python-package/xgboost/spark/data.py
@@ -321,6 +321,10 @@ def split_params() -> Tuple[Dict[str, Any], Dict[str, Union[int, float, bool]]]:
     end = time.time()
     get_logger("XGBoost-PySpark").warning("cache_partitions takes {} seconds".format(round(end - start, 2)))
 
+    from pyspark import BarrierTaskContext
+    context = BarrierTaskContext.get()
+    context.barrier()
+
     start = time.time()
     if feature_cols is not None and use_qdm:
         dtrain: DMatrix = make_qdm(train_data, dev_ordinal, meta, None, params)