Avoid copying an AMRM delegation token from driver to AM in deploy-mo…

…de client An AMRM delegation token can exist when the spark submit is issue from another hadoop container (example of how cuttle launches jobs on hadoop) --------- Co-authored-by: w.montaz <[email protected]>
criteo-forks · Jun 14, 2024 · 18cd536 · 18cd536
1 parent 1f247d4
commit 18cd536
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala b/core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala
@@ -35,6 +35,7 @@ import org.apache.hadoop.mapred.JobConf
 import org.apache.hadoop.security.{Credentials, UserGroupInformation}
 import org.apache.hadoop.security.token.{Token, TokenIdentifier}
 import org.apache.hadoop.security.token.delegation.AbstractDelegationTokenIdentifier
+import org.apache.hadoop.yarn.security.AMRMTokenIdentifier
 
 import org.apache.spark.{SparkConf, SparkException}
 import org.apache.spark.internal.Logging
@@ -154,9 +155,20 @@ private[spark] class SparkHadoopUtil extends Logging {
   private[spark] def addDelegationTokens(tokens: Array[Byte], sparkConf: SparkConf): Unit = {
     UserGroupInformation.setConfiguration(newConfiguration(sparkConf))
     val creds = deserialize(tokens)
+
+    val filteredCreds = new Credentials()
+    for ((k, v) <- creds.getSecretKeyMap.asScala) {
+      filteredCreds.addSecretKey(k, v)
+    }
+    for ((k, token) <- creds.getTokenMap.asScala) {
+      if (!token.getKind.equals(AMRMTokenIdentifier.KIND_NAME)) {
+        filteredCreds.addToken(k, token)
+      }
+    }
+
     logInfo("Updating delegation tokens for current user.")
-    logDebug(s"Adding/updating delegation tokens ${dumpTokens(creds)}")
-    addCurrentUserCredentials(creds)
+    logDebug(s"Adding/updating delegation tokens ${dumpTokens(filteredCreds)}")
+    addCurrentUserCredentials(filteredCreds)
   }
 
   /**