[BUG] Does spark-rapids-ml support "unbalanced" configurations?

[an-ys]

In the latest reply in #853, I mentioned that I faced a weird issue regarding some overhead for some applications.

It seems like for KMeans and PCA, Spark is assigning the tasks to the wrong executors in standalone mode. This happens when I use five GPUs (i.e., 4 GPUs on one node, node A, and 1 GPU on another node, node B), but it doesn't seem to happen when I use four GPUs. This seems to happen all the time for smaller sizes, but for other sizes, there are some cases when it is faster than running with GPUs when running the benchmark multiple times. Do note that my parquet is divided into ~30MiB files, but I don't think it should matter much.

When running KMeans on a cluster of two nodes, with one node having 4 GPUs and the other 1 GPU, it tries to run most of the tasks in only one executor (the lone GPU in the other node). It runs some tasks in the node with four GPUs, but it only uses one GPU there. As a result, instead of having node-level/process-level tasks in other stages, the application has tasks with the locality level of "any", and it ends up running the tasks somewhat sequentially since the maximum number of concurrent tasks for one GPU is 2.

When running PCA on 6 GPUs (either with 2 unbalanced nodes or 3 nodes of 2 gpus), it runs all the tasks in only one node (the one with 2 GPUs) for some stages instead of all 6 executors. This can be seen below. If you check the second image, you would see that most of the work is done by two executors in a single node (IP address ends with 101).

As mentioned in the second paragraph, there are some cases where it assigns the tasks correctly, but it does not most of the time. This picture below shows what happens when Spark assigns the tasks correctly or incorrectly. The slowest configuration is about 2 times slower than the fastest configuration since Spark has to wait for the first two tasks to finish before assigning the last two tasks as the max # of concurrent tasks is set to 2.

(Also, the directory is called spark-rapids-ml-24.08-benchmark, but I have been updating the files to match the ones in the repo. I used Spark RAPIDS 25.02 and Spark RAPIDS ML 25.02.)

This does not seem to happen in other workloads (i.e., linear regression, logistic regression, and random forest classifier).

[an-ys]

Also, I am facing another issue that I was originally going to add to #786 as a reply, but I think it is more related to this issue.

Running the KMeans with 10 GPUs seems to make the application hang for no apparent reason. After 10 hours, it is still stuck at the fit stage. The data size does not seem to matter (4GB or 32GB). I am using 4 nodes with two nodes having one GPU each and the other two nodes having four GPUs each don't recall if this issue also occurs for a cluster of 3 nodes with a total of 10 GPUs).

I tried to check if it is related to this issue, and it seems to be the case. When checking stage 2 [$anonfun$internalDoExecuteColumnar$1 at Spark320PlusShims.scala:339], which is "GpuShuffleExchangeExecBase", four tasks are assigned to the node with one GPU while the other six tasks are assigned correctly to two nodes (one node has 4 tasks on 4 GPUs, and the other has two tasks on 2 GPU), but the locality level of these tasks is set to ANY instead of NODE_LOCAL, which is the level the first four tasks have.

For stage 7, the fit stage where the application is stuck, a task is exclusively assigned, which is to be expected

Stage 2

Stage 7 (stuck)

[eordentlich]

Is your standalone cluster showing expected resources? Do you have 4 workers one for each node and the respective number of GPUs showing for each worker on the UI (http://<master_host>:8080). Can you tell where the executors are hanging? Either stack trace or stderr logs might give a clue?

[an-ys]

Each node has 4 GPUs, but I restrict the number of workers to a specific configuration.

For example, one of the configurations I used for 10 GPUs consists of two nodes with 4 GPUs and two nodes with a single GPU as seen here.

I also tried using a different configuration for 10 GPUs, which is a cluster of three nodes with two nodes having four GPUs and the other having two GPUs, and it works.

For 11 GPUs, one of the nodes use two GPUs instead of just one. The workers are dead here as I ran KMeans with 10 GPUs afterward. This configuration also works for KMeans.

spark_rapids_ml.clustering.KMeans - INFO - Loading data into python worker memory
spark_rapids_ml.clustering.KMeans - INFO - Initializing cuml context
INFO BarrierTaskContext: Task n from Stage 7(Attempt 0) has entered the global sync, current barrier epoch is 0.
INFO BarrierTaskContext: Task n from Stage 7(Attempt 0) finished global sync successfully, waited for n seconds, current barrier epoch is 1.
NCCL version 2.25.1+cuda12.8
spark_rapids_ml.clustering.KMeans - INFO - Invoking cuml fit
107526000
...../spark-rapids-ml-benchmark/.venv/lib/python3.12/site-packages/cuml/internals/api_decorators.py:193: FutureWarning: The default value of `n_init` will change from 1 to 'auto' in 25.04. Set the value of `n_init` explicitly to suppress this warning.
  ret = func(*args, **kwargs)

NCCL version is not printed for one of the nodes. I tried downgrading the NCCL to 2.24.3-1 since I use 12.6 instead of 12.8, and it no longer prints the NCCL version as seen below, but it still gets stuck. It also prints a number (i.e., 107517000, 107523000, 107526000, 107529000, 107535000, 107538000). I am guessing it is a version incompatibility issue?

spark_rapids_ml.clustering.KMeans - INFO - Loading data into python worker memory
spark_rapids_ml.clustering.KMeans - INFO - Initializing cuml context
INFO BarrierTaskContext: Task 58 from Stage 7(Attempt 0) has entered the global sync, current barrier epoch is 0.
INFO BarrierTaskContext: Task 58 from Stage 7(Attempt 0) finished global sync successfully, waited for 3 seconds, current barrier epoch is 1.
spark_rapids_ml.clustering.KMeans - INFO - Invoking cuml fit
<some number>
...../spark-rapids-ml-benchmark/.venv/lib/python3.12/site-packages/cuml/internals/api_decorators.py:193: FutureWarning: The default value of `n_init` will change from 1 to 'auto' in 25.04. Set the value of `n_init` explicitly to suppress this warning.
  ret = func(*args, **kwargs)

The thing is that this only happens for this specific configuration and not for other configurations. When I run KMeans with 10 GPUs but with a different configuration (with 3 nodes instead of 4), it works without any issues.

spark_rapids_ml.clustering.KMeans - INFO - Loading data into python worker memory
spark_rapids_ml.clustering.KMeans - INFO - Initializing cuml context
INFO BarrierTaskContext: Task n from Stage 7(Attempt 0) has entered the global sync, current barrier epoch is 0.
INFO BarrierTaskContext: Task n from Stage 7(Attempt 0) finished global sync successfully, waited for 1 seconds, current barrier epoch is 1.
spark_rapids_ml.clustering.KMeans - INFO - Invoking cuml fit
107538000
...../spark-rapids-ml-benchmark/.venv/lib/python3.12/site-packages/cuml/internals/api_decorators.py:193: FutureWarning: The default value of `n_init` will change from 1 to 'auto' in 25.04. Set the value of `n_init` explicitly to suppress this warning.
  ret = func(*args, **kwargs)
spark_rapids_ml.clustering.KMeans - INFO - iterations: 31, inertia: 0.0
spark_rapids_ml.clustering.KMeans - INFO - Cuml fit complete
INFO BarrierTaskContext: Task n from Stage 7(Attempt 0) has entered the global sync, current barrier epoch is 1.
INFO BarrierTaskContext: Task n from Stage 7(Attempt 0) finished global sync successfully, waited for 1 seconds, current barrier epoch is 2.
INFO GpuMapInBatchExec$$anon$1: Times: total = 19259, boot = 687, init = 8428, finish = 10144
INFO PythonRunner: Times: total = 18374, boot = 725, init = 17649, finish = 0
INFO Executor: Finished task 7.0 in stage 7.0 (TID 58). 6381 bytes result sent to driver
INFO CoarseGrainedExecutorBackend: Got assigned task 65
INFO Executor: Running task 7.0 in stage 9.0 (TID 65)
INFO TorrentBroadcast: Started reading broadcast variable 13 with 1 pieces (estimated total size 4.0 MiB)
INFO MemoryStore: Block broadcast_13_piece0 stored as bytes in memory (estimated size 23.9 KiB, free 29.4 GiB)
INFO TorrentBroadcast: Reading broadcast variable 13 took 13 ms
INFO MemoryStore: Block broadcast_13 stored as values in memory (estimated size 53.4 KiB, free 29.4 GiB)
INFO BlockManager: Found block rdd_10_7 locally
INFO TorrentBroadcast: Started reading broadcast variable 12 with 1 pieces (estimated total size 4.0 MiB)
INFO MemoryStore: Block broadcast_12_piece0 stored as bytes in memory (estimated size 3.0 KiB, free 29.4 GiB)
INFO TorrentBroadcast: Reading broadcast variable 12 took 11 ms
INFO MemoryStore: Block broadcast_12 stored as values in memory (estimated size 24.3 KiB, free 29.4 GiB)
INFO TorrentBroadcast: Started reading broadcast variable 11 with 4 pieces (estimated total size 16.0 MiB)
INFO MemoryStore: Block broadcast_11_piece1 stored as bytes in memory (estimated size 4.0 MiB, free 29.4 GiB)
INFO MemoryStore: Block broadcast_11_piece3 stored as bytes in memory (estimated size 3.8 MiB, free 29.4 GiB)
INFO MemoryStore: Block broadcast_11_piece2 stored as bytes in memory (estimated size 4.0 MiB, free 29.4 GiB)
INFO MemoryStore: Block broadcast_11_piece0 stored as bytes in memory (estimated size 4.0 MiB, free 29.4 GiB)
INFO TorrentBroadcast: Reading broadcast variable 11 took 120 ms
25/05/09 06:45:13 INFO MemoryStore: Block broadcast_11 stored as values in memory (estimated size 312.0 B, free 29.4 GiB)
...../spark-rapids-ml-benchmark/.venv/lib/python3.12/site-packages/cuml/internals/api_decorators.py:193: FutureWarning: The default value of `n_init` will change from 1 to 'auto' in 25.04. Set the value of `n_init` explicitly to suppress this warning.
  ret = func(*args, **kwargs)

The returned number also differs between different executors for the successful configuration with 10 GPUs, so I don't think that is the issue. The other configuration just gets stuck at the cuML fit stage without logging anything. I also tried using at least 20 different configurations, and it is only this configuration that leads to an infinite loop for KMeans. More configurations have the first issue, which is assigning tasks to the wrong executor/GPU.

Also, running with all 16 GPUs works fine, so the problem is not with the GPUs themselves. I have also tried changing the GPUs used for the application as seen on the figure, where the GPU order starts from the last GPU ID instead of the first.