debug CI2

dmlc · Jul 29, 2022 · a3b1060 · a3b1060
1 parent 4a39302
commit a3b1060
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/python-package/xgboost/spark/data.py b/python-package/xgboost/spark/data.py
@@ -82,9 +82,7 @@ def _fetch(self, data: Optional[Sequence[pd.DataFrame]]) -> Optional[pd.DataFram
             # See https://github.com/rapidsai/cudf/issues/11386
             cp.cuda.runtime.setDevice(self._device_id)
             gpu_id = cp.cuda.runtime.getDevice()
-            import pyspark
-            pid = pyspark.TaskContext.get().partitionId()
-            print(f"partiton id {pid} _fetch got gpu id {gpu_id}")
+            print(f" _fetch got gpu id {gpu_id}, device_id: {self._device_id}")
             return cudf.DataFrame(data[self._iter])
 
         return data[self._iter]