handle NumberFormatException with executorId (#76)

IBMSparkGPU · Dec 13, 2017 · 46b5dd5 · 46b5dd5
1 parent 3fb5218
commit 46b5dd5
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/gpu-enabler/src/main/scala/com/ibm/gpuenabler/GPUSparkEnv.scala b/gpu-enabler/src/main/scala/com/ibm/gpuenabler/GPUSparkEnv.scala
@@ -90,7 +90,15 @@ private[gpuenabler] object GPUSparkEnv {
           if (env.isGPUEnabled) {
             val executorId = env.executorId match {
               case "driver" => 0
-              case _ => SparkEnv.get.executorId.toInt
+              case _ => 
+                val execIdStr = SparkEnv.get.executorId
+                if (execIdStr.forall( _.isDigit )) {
+                  execIdStr.substring(0, Math.min(execIdStr.length, 5)).toInt
+                } else {
+                  val stripNum = ("[0-9]".r findAllIn execIdStr).mkString("")
+                  if (stripNum.isEmpty) 0
+                  else stripNum.substring(0, Math.min(stripNum.length, 5)).toInt
+                }
             }
 
             env.gpuDevice = executorId % env.gpuCount