On task failure catch some CUDA exceptions and kill executor [databricks]

[tgravescs]

Related to #5029. This is shorter term solution to just parse the exception message to catch certain types of unrecoverable CUDA errors. It may not be bullet proof as the messages could change.

Here if we find an exception that we think is unrecoverable we system.exit to kill the executor.
Generally you would want to use this with the Spark excludeOnFailure functionality so it doesn't start the executor back up using the same GPU.

I've manually tested this by faking the exception occurring since we can't reproduce it. It properly kills the executor when it sees the exception.

Sample code used to cause failures:

sc.range(0, 4, 1, 4).mapPartitions{x =>
  import org.apache.spark.TaskContext
  val tc = TaskContext.get()
  println("task id: " + tc.taskAttemptId)
  
  if (tc.taskAttemptId % 3 == 0) {
    try {
    throw new Exception("cudaErrorHardwareStackError")
    } catch {
      case e: Throwable =>
      throw new RuntimeException(s"CUDA error encountered: ${e.getMessage}", e)
    }
  }
  x.map(x => x)}.collect()

this generates exceptions like:

22/03/31 22:46:57 ERROR Executor: Exception in task 0.0 in stage 0.0 (TID 0)
java.lang.RuntimeException: CUDA error encountered: cudaErrorHardwareStackError

This new executor plugin code catches and logs the following and then exits:

22/03/31 22:46:57 ERROR RapidsExecutorPlugin: Stopping the Executor based on exception being a fatal CUDA error: java.lang.RuntimeException: CUDA error encountered: cudaErrorHardwareStackError

In standalone mode with the excludeOnFailure spark configs set to 1 for the node exclusion, it when the task fails and this kills the executor, the node with be excluded and the worker will not be able to restart an executor on that node. Also note keep in mind the spark config spark.excludeOnFailure.timeout which will try spark to retry that node after the timeout value.
Without excludeonFailure, the executors just get restarted on the same nodes for standalone mode. I tested on yarn as well and there it will restart executors but it could be on different nodes depending on the size of the cluster.

[tgravescs]

build

[tgravescs]

build