xgboost4j-spark-gpu fails to run on multi-GPU server with GPUs in exclusive process mode and spark-rapids plugin

[eordentlich]

Exceptions result when running as above for both driver and executor processes. Stack traces for both includes
/workspace/src/common/cuda_rt_utils.cc: 60: cudaErrorDevicesUnavailable: CUDA-capable device(s) is/are busy or unavailable
indicating XGBoost code is attempting to create CUDA contexts even though the corresponding Spark driver and executor processes should either not be using GPUs (Spark driver) or have been assigned other device ids by Spark.

On the driver side, it seems one source (might be more than one with the first one triggering the error) might be coming from RabitTracker initialization which uses the common InitNewThread function whose logic seems to setCurrentDevice on the default device. The source on the executor side is not clear at this point, but it is XGBoost code.

When switching GPUs to default process mode, I believe this behavior (depending on root cause above for executors, i.e. is it coming from main executor thread or new thread) can lead to race conditions in current stream/thread device settings with spark-rapids, resulting in spurious cuda kernels launched by executors on non-assigned device 0, in turn causing errors like:

Caused by: ai.rapids.cudf.CudfException: after determining tmp storage requirements for exclusive_scan: cudaErrorInvalidDevice: invalid device ordinal

first noted in NVIDIA/spark-rapids-examples#565

Specifically, executor processes launching kernels on assigned device and device 0 can be observed in nsys traces.

[trivialfis]

Thank you for sharing the issue. Do you intend to use XGBoost with GPUs in your current setting? Or do you prefer XGBoost to not touch any CUDA API and exclusively use CPU?

[eordentlich]

Intention is to use GPUs on executors, but not on the driver. Actually was working on root causing NVIDIA/spark-rapids-examples#565 . After observing unexpected device usage by executor processes attempted to get more insight using exclusive process mode and encountered the noted behavior, which was unexpected as executor processes should not be changing devices when spark-rapids is enabled. Seems the spark driver should also not be creating spurious CUDA contexts.

[trivialfis]

Got it. XGBoost does try to obtain information like the current device when calling its API. We can probably map a code path for the tracker to not touch the CUDA runtime. But that will take some very careful tests, I'm not sure how to do it consistently in the long run.

[eordentlich]

I could be wrong, but cudaGetDevice and such should be ok to get info. cudaSetDevice is what could be problematic (i.e. creates the CUDA context). Unless a use case requires exclusive mode, the tracker invoking cudaSetDevice on the driver is just incidental in this case and may use some (minimal) resources unnecessarily, but the executor process side may be leading to inconsistent device setting wrt to spark-rapids rmm pooling resources and kernel launches and such, leading to instability and errors. However, I'm not sure if that is in the worker side of the tracker/communicator or has another source at this point.

[trivialfis]

Based on https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html#group__CUDART__DEVICE_1g80861db2ce7c29b6e8055af8ae01bc78
Excerpt:

Note that this function may also return cudaErrorInitializationError, cudaErrorInsufficientDriver or cudaErrorNoDevice if this call tries to initialize internal CUDA RT state. 

The get device call can initialize the CUDA context.

[trivialfis]

In XGBoost, we let some of the CUDA calls to fail, like the get device function, so that XGBoost can run on CPU-only nodes even if compiled with CUDA support.
I'm not sure why exclusive mode is different from a CPU-only machine from XGBoost's perspective. Will need to do some digging.

[eordentlich]

This patch to set the device in the main Spark train mapPartitions thread seems to fix the issue from the executor side. Without this, InitNewThread just propagates the default device 0 of this initial thread to other child Communicator threads. The driver issue aside (which can be addressed by setting CUDA_VISIBLE_DEVICES="" when launching Spark), this fix allows running in exclusive process mode.
It also seems to resolve the spark-rapids-examples issue in limited testing. Not fully understanding this since the spurious device 0 access from XGBoost seem to be introduced on the workers during training by the various OS Communicator thread(s), while the error in the example arises in OS native threads associated with jvm threads. One would not think there would be such native thread reuse by the jvm (is it possible with std::async?).

The patch may not be ideal since it requires the cudf jni package as a dependency even in the case of gpu accelerated Spark without the spark-rapids plugin. With the spark-rapids plugin this is fine since cudf jni is bundled with the spark-rapids jar. Perhaps @wbo4958 can take a pass. One option is adding a jni for cudaSetDevices to XGBoost jni?

[wbo4958]

Thx @eordentlich, I just read the code. Were you trying to say, do we need to set the device before creating the Communicator where it's going to use default device 0 before using the correct device ID got from Spark?

Hi @trivialfis, may I ask will XGCommunicatorInit use default device ID 0 for the GPU pipeline? If in that case, Can we pass the device ID from Spark as a DMLC variable just like (DMLC_TASK_ID) to Communicator when initializing?

[trivialfis]

The communicator is not used on the driver. Only the tracker is initialized.

I think we need an actual setup for testing before we can debug it. I don't have a spark cluster with exclusive access mode, please let me know if you can configure one for replicating the issue.

[eordentlich]

Thx @eordentlich, I just read the code. Were you trying to say, do we need to set the device before creating the Communicator where it's going to use default device 0 before using the correct device ID got from Spark?

Yes and it is not just the Communicator threads. The XGBoost C-api saves and restores the default device on each api invocation (via XGBoostAPIGuard) leading to many (> 1000 in one case) calls to cudaSetDevice where it is just restoring back to the original default device. If that device is not set in the calling thread, it will use device 0. I believe NCCL also does something similar. So ideally, the device assigned by Spark is set as the default device on any thread that might invoke XGBoost. Spark-rapids plugin jni actually does this on every api invocation, but not XGBoost4j-spark currently.

[wbo4958]

That's great. I can repro this issue locally, let me dig it first.

[trivialfis]

I will try to see if I can make the set device entirely optional in XGBoost. Afterward, we can set the device ID upfront for the Spark interface.

[wbo4958]

I checked Spark-rapids and CUDF, looks like both of them just require to setDevice only once globally at very beginning.