update

wbo4958 · Sep 4, 2023 · 1e97ae8 · 1e97ae8
1 parent d11e55a
commit 1e97ae8
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/python-package/xgboost/spark/core.py b/python-package/xgboost/spark/core.py
@@ -1256,17 +1256,19 @@ def predict_udf(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.Series]:
             dev_ordinal = -1
             if is_local:
                 if run_on_gpu and is_cupy_available():
-                    total_gpus = cupy.cuda.runtime.getDeviceCount()
+                    import cupy as cp  # pylint: disable=import-error
+
+                    total_gpus = cp.cuda.runtime.getDeviceCount()
                     if total_gpus > 0:
-                        partition_id = context.partitionId()
+                        from pyspark import TaskContext
+                        partition_id = TaskContext.get().partitionId()
                         # For transform local mode, default the gpu_id to (partition id) % gpus.
                         dev_ordinal = partition_id % total_gpus
             elif run_on_gpu:
                 from pyspark import TaskContext
                 dev_ordinal = _get_gpu_id(TaskContext.get())
 
             if dev_ordinal >= 0:
-                print("------------------- ", TaskContext.get().partitionId())
                 device = "cuda:" + str(dev_ordinal)
                 get_logger("XGBoost-PySpark").info(
                     "Do the inference with device: %s", device
@@ -1278,7 +1280,7 @@ def predict_udf(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.Series]:
                     )
 
             def to_gpu(data: ArrayLike) -> ArrayLike:
-                """Move the data to gpu"""
+                """Move the data to gpu if possible"""
                 if dev_ordinal >= 0:
                     import cudf  # pylint: disable=import-error
                     import cupy as cp  # pylint: disable=import-error