From 3ad097a54bf5ed1676e9604cb18ae6c6fca6e747 Mon Sep 17 00:00:00 2001 From: Kousuke Saruta Date: Wed, 29 Jul 2026 18:10:35 +0900 Subject: [PATCH] [SPARK-57894][SPARK-57895][CORE] UpdateUserCredentials RPC, backend integration, and initial credential delivery Add the credential distribution layer for the OIDC Credential Propagation SPIP (SPARK-57703). This covers Sub-task 5 (UpdateUserCredentials RPC and backend integration) and Sub-task 6 (initial credential delivery for newly registered executors). Credential delivery via three complementary paths: 1. SparkAppConfig - initial delivery on executor registration 2. UpdateUserCredentials RPC - broadcast on credential renewal 3. TaskDescription - per-task delivery to eliminate race conditions Key implementation details: - UpdateUserCredentials(credentials: Array[Byte]) new RPC message - SparkEnv.userCredentials: AtomicReference credential store - CoarseGrainedSchedulerBackend: starts/stops UserCredentialManager, handles UpdateUserCredentials in DriverEndpoint.receive - CoarseGrainedExecutorBackend: handles UpdateUserCredentials, applies initial credentials from SparkAppConfig - TaskDescription: new userCredentials field with encode/decode - Executor: applies credentials before task execution Security: raw identity tokens (UserContext.rawToken) are never present in any RPC payload, TaskDescription, or SparkAppConfig. Only derived ServiceCredential bundles (serialized UserCredentials) are transmitted. Closes SPARK-57894, SPARK-57895 --- .../scala/org/apache/spark/SparkEnv.scala | 10 +++ .../CoarseGrainedExecutorBackend.scala | 11 +++ .../org/apache/spark/executor/Executor.scala | 9 +++ .../spark/scheduler/TaskDescription.scala | 23 +++++- .../spark/scheduler/TaskSetManager.scala | 1 + .../cluster/CoarseGrainedClusterMessage.scala | 4 + .../CoarseGrainedSchedulerBackend.scala | 54 +++++++++++++ .../CoarseGrainedExecutorBackendSuite.scala | 6 +- .../apache/spark/executor/ExecutorSuite.scala | 1 + .../CoarseGrainedSchedulerBackendSuite.scala | 79 ++++++++++++++++++- .../scheduler/TaskDescriptionSuite.scala | 64 +++++++++++++++ .../k8s/KubernetesExecutorBackend.scala | 5 ++ 12 files changed, 260 insertions(+), 7 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala index ca48ee473eb0f..557807f345458 100644 --- a/core/src/main/scala/org/apache/spark/SparkEnv.scala +++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala @@ -19,6 +19,7 @@ package org.apache.spark import java.io.File import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference import scala.collection.concurrent import scala.collection.mutable @@ -268,6 +269,15 @@ class SparkEnv ( private[spark] var executorBackend: Option[ExecutorBackend] = None + /** + * Credential store for OIDC-based user credentials on executors. + * Updated via `UpdateUserCredentials` RPC and `TaskDescription` credential delivery. + * Read by connector-specific credential providers (e.g., SparkOidcAwsCredentialsProvider). + * Contains serialized `UserCredentials` (no raw identity token). + */ + private[spark] val userCredentials: AtomicReference[Array[Byte]] = + new AtomicReference[Array[Byte]]() + private[spark] def stop(): Unit = { if (!isStopped) { diff --git a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala index 206a6a0fe385c..175a5ea94deea 100644 --- a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala +++ b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala @@ -223,6 +223,11 @@ private[spark] class CoarseGrainedExecutorBackend( logInfo(log"Received tokens of ${MDC(LogKeys.NUM_BYTES, tokenBytes.length)} bytes") SparkHadoopUtil.get.addDelegationTokens(tokenBytes, env.conf) + case UpdateUserCredentials(credentials) => + logInfo(log"Received user credentials of " + + log"${MDC(LogKeys.NUM_BYTES, credentials.length)} bytes") + env.userCredentials.set(credentials) + case DecommissionExecutor => decommissionSelf() } @@ -505,6 +510,12 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging { } val env = SparkEnv.createExecutorEnv(driverConf, arguments.executorId, arguments.bindAddress, arguments.hostname, arguments.cores, cfg.ioEncryptionKey, isLocal = false) + + // Apply initial user credentials to the executor credential store. + cfg.userCredentials.foreach { credentials => + env.userCredentials.set(credentials) + } + // Set the application attemptId in the BlockStoreClient if available. val appAttemptId = env.conf.get(APP_ATTEMPT_ID) appAttemptId.foreach(attemptId => diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala b/core/src/main/scala/org/apache/spark/executor/Executor.scala index f288ceacef7e6..be2c17d537767 100644 --- a/core/src/main/scala/org/apache/spark/executor/Executor.scala +++ b/core/src/main/scala/org/apache/spark/executor/Executor.scala @@ -847,6 +847,15 @@ private[spark] class Executor( // requires access to properties contained within (e.g. for access control). Executor.taskDeserializationProps.set(taskDescription.properties) + // Apply user credentials from TaskDescription to the executor credential store. + // This ensures credentials are available before any task code runs, avoiding + // the race between RPC broadcast and task dispatch. + // All tasks in a Spark application belong to the same user, so last-writer-wins + // is acceptable when multiple task threads update concurrently. + taskDescription.userCredentials.foreach { creds => + env.userCredentials.set(creds) + } + updateDependencies( taskDescription.artifacts.files, taskDescription.artifacts.jars, diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala index 7eba116c56800..12fc314a66dfe 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala @@ -61,6 +61,7 @@ private[spark] class TaskDescription( // Eg, Map("gpu" -> Map("0" -> ResourceAmountUtils.toInternalResource(0.7))): // assign 0.7 of the gpu address "0" to this task val resources: immutable.Map[String, immutable.Map[String, Long]], + val userCredentials: Option[Array[Byte]], val serializedTask: ByteBuffer) { assert(cpus > 0, "CPUs per task should be > 0") @@ -121,6 +122,16 @@ private[spark] object TaskDescription { // Write resources. serializeResources(taskDescription.resources, dataOut) + // Write user credentials (OIDC). + taskDescription.userCredentials match { + case Some(creds) => + dataOut.writeBoolean(true) + dataOut.writeInt(creds.length) + dataOut.write(creds) + case None => + dataOut.writeBoolean(false) + } + // Write the task. The task is already serialized, so write it directly to the byte buffer. Utils.writeByteBuffer(taskDescription.serializedTask, bytesOut) @@ -228,10 +239,20 @@ private[spark] object TaskDescription { // Read resources. val resources = deserializeResources(dataIn) + // Read user credentials (OIDC). + val userCredentials = if (dataIn.readBoolean()) { + val length = dataIn.readInt() + val creds = new Array[Byte](length) + dataIn.readFully(creds) + Some(creds) + } else { + None + } + // Create a sub-buffer for the serialized task into its own buffer (to be deserialized later). val serializedTask = byteBuffer.slice() new TaskDescription(taskId, attemptNumber, executorId, name, index, partitionId, artifacts, - properties, cpus, resources, serializedTask) + properties, cpus, resources, userCredentials, serializedTask) } } diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala index 4fbe197cb8382..57ecfd6771556 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -604,6 +604,7 @@ private[spark] class TaskSetManager( task.localProperties, taskCpus, taskResourceAssignments, + Option(SparkEnv.get.userCredentials.get()), serializedTask) } diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala index 84a6f9b3be2a1..bc6016b79ba2e 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala @@ -35,6 +35,7 @@ private[spark] object CoarseGrainedClusterMessages { sparkProperties: Seq[(String, String)], ioEncryptionKey: Option[Array[Byte]], hadoopDelegationCreds: Option[Array[Byte]], + userCredentials: Option[Array[Byte]], resourceProfile: ResourceProfile, logLevel: Option[String]) extends CoarseGrainedClusterMessage @@ -60,6 +61,9 @@ private[spark] object CoarseGrainedClusterMessages { case class UpdateDelegationTokens(tokens: Array[Byte]) extends CoarseGrainedClusterMessage + case class UpdateUserCredentials(credentials: Array[Byte]) + extends CoarseGrainedClusterMessage + // Executors to driver case class RegisterExecutor( executorId: String, diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala index 99fdc534a6fb9..acad934eeba5f 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala @@ -28,6 +28,7 @@ import com.google.common.cache.CacheBuilder import org.apache.spark.{ExecutorAllocationClient, SparkEnv, TaskState} import org.apache.spark.deploy.SparkHadoopUtil +import org.apache.spark.deploy.security.UserCredentialManager import org.apache.spark.errors.SparkCoreErrors import org.apache.spark.executor.ExecutorLogUrlHandler import org.apache.spark.internal.{config, Logging} @@ -132,6 +133,13 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp // Current set of delegation tokens to send to executors. private val delegationTokens = new AtomicReference[Array[Byte]]() + // Current set of OIDC user credentials to send to executors. + // Updated by UserCredentialManager and sent via UpdateUserCredentials RPC. + private val userCredentials = new AtomicReference[Array[Byte]]() + + // UserCredentialManager for OIDC credential propagation (if enabled). + private var userCredentialManager: Option[UserCredentialManager] = None + private val reviveThread = ThreadUtils.newDaemonSingleThreadScheduledExecutor("driver-revive-thread") @@ -220,6 +228,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp case UpdateDelegationTokens(newDelegationTokens) => updateDelegationTokens(newDelegationTokens) + case UpdateUserCredentials(newCredentials) => + updateUserCredentials(newCredentials) + case RemoveExecutor(executorId, reason) => // We will remove the executor's state and cannot restore it. However, the connection // between the driver and the executor may be still alive so that the executor won't exit @@ -358,6 +369,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp sparkProperties, SparkEnv.get.securityManager.getIOEncryptionKey(), Option(delegationTokens.get()), + Option(userCredentials.get()), rp, currentLogLevel) context.reply(reply) @@ -618,6 +630,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp override def start(): Unit = { setupTokenManager() + setupUserCredentialManager() if (conf.get(DIRECT_CREDENTIAL_PROVIDERS_ENABLED) && delegationTokenManager.isEmpty) { logWarning("spark.security.directCredentialProviders.enabled is set but " + "this cluster manager does not support credential distribution. " + @@ -648,6 +661,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp cleanupService.foreach(_.shutdownNow()) stopExecutors() stopTokenManager() + stopUserCredentialManager() try { if (driverEndpoint != null) { driverEndpoint.askSync[Boolean](StopDriver) @@ -1047,6 +1061,46 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp protected def currentDelegationTokens: Array[Byte] = delegationTokens.get() + /** + * Update user credentials and broadcast to all registered executors. + * Called from the DriverEndpoint receive loop (thread-safe access to executorDataMap). + */ + private def updateUserCredentials(credentials: Array[Byte]): Unit = { + userCredentials.set(credentials) + // Also update driver's SparkEnv so TaskDescription can include current credentials. + SparkEnv.get.userCredentials.set(credentials) + executorDataMap.values.foreach { ed => + ed.executorEndpoint.send(UpdateUserCredentials(credentials)) + } + } + + protected def currentUserCredentials: Array[Byte] = userCredentials.get() + + /** + * Start the UserCredentialManager if OIDC credential propagation is enabled. + * Called from start(), independently of Kerberos/HadoopDelegationTokenManager. + */ + private def setupUserCredentialManager(): Unit = { + userCredentialManager = UserCredentialManager.create(conf, { credentials => + // Send to DriverEndpoint to ensure thread-safe access to executorDataMap. + // This mirrors HadoopDelegationTokenManager's pattern of sending + // UpdateDelegationTokens via schedulerRef. + driverEndpoint.send(UpdateUserCredentials(credentials)) + }) + userCredentialManager.foreach { manager => + val initialCredentials = manager.start() + // Store initial credentials for SparkAppConfig (late-registering executors) + // and TaskDescription (task dispatch). Both stores must be populated synchronously + // to avoid a window where tasks could be dispatched without credentials. + userCredentials.set(initialCredentials) + SparkEnv.get.userCredentials.set(initialCredentials) + } + } + + private def stopUserCredentialManager(): Unit = { + userCredentialManager.foreach(_.stop()) + } + /** * Checks whether the executor is excluded due to failure(s). This is called when the executor * tries to register with the scheduler, and will deny registration if this method returns true. diff --git a/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala b/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala index a038d8e8613ef..c7307066c78d3 100644 --- a/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala +++ b/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala @@ -312,7 +312,7 @@ class CoarseGrainedExecutorBackendSuite extends SparkFunSuite // We don't really verify the data, just pass it around. val data = ByteBuffer.wrap(Array[Byte](1, 2, 3, 4)) val taskDescription = new TaskDescription(taskId, 2, "1", "TASK 1000000", 19, - 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, data) + 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, None, data) val serializedTaskDescription = TaskDescription.encode(taskDescription) backend.rpcEnv.setupEndpoint("Executor 1", backend) backend.executor = mock[Executor](CALLS_REAL_METHODS) @@ -437,7 +437,7 @@ class CoarseGrainedExecutorBackendSuite extends SparkFunSuite // Fake tasks with different taskIds. val taskDescriptions = (1 to numTasks).map { taskId => new TaskDescription(taskId, 2, "1", s"TASK $taskId", 19, - 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, data) + 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, None, data) } assert(taskDescriptions.length == numTasks) @@ -531,7 +531,7 @@ class CoarseGrainedExecutorBackendSuite extends SparkFunSuite // Fake tasks with different taskIds. val taskDescriptions = (1 to numTasks).map { taskId => new TaskDescription(taskId, 2, "1", s"TASK $taskId", 19, - 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, data) + 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, None, data) } assert(taskDescriptions.length == numTasks) diff --git a/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala b/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala index 4362eb2f816f1..80b0bf69d45cf 100644 --- a/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala +++ b/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala @@ -870,6 +870,7 @@ class ExecutorSuite extends SparkFunSuite properties = new Properties, cpus = 1, resources = Map.empty, + None, serializedTask) } diff --git a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala index 99b757496379d..dcd9fe4348586 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala @@ -348,7 +348,7 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo val taskCpus = 1 val taskDescs: Seq[Seq[TaskDescription]] = Seq(Seq(new TaskDescription(1, 0, "1", "t1", 0, 1, JobArtifactSet.emptyJobArtifactSet, new Properties(), - taskCpus, taskResources, bytebuffer))) + taskCpus, taskResources, None, bytebuffer))) val ts = backend.getTaskSchedulerImpl() when(ts.resourceOffers(any[IndexedSeq[WorkerOffer]], any[Boolean])).thenReturn(taskDescs) @@ -455,7 +455,7 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo val taskCpus = 1 val taskDescs: Seq[Seq[TaskDescription]] = Seq(Seq(new TaskDescription(1, 0, "1", "t1", 0, 1, JobArtifactSet.emptyJobArtifactSet, new Properties(), - taskCpus, taskResources, bytebuffer))) + taskCpus, taskResources, None, bytebuffer))) val ts = backend.getTaskSchedulerImpl() when(ts.resourceOffers(any[IndexedSeq[WorkerOffer]], any[Boolean])).thenReturn(taskDescs) @@ -548,7 +548,7 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo val taskCpus = 2 val taskDescs: Seq[Seq[TaskDescription]] = Seq(Seq(new TaskDescription(1, 0, "1", "t1", 0, 1, JobArtifactSet.emptyJobArtifactSet, new Properties(), - taskCpus, Map.empty, bytebuffer))) + taskCpus, Map.empty, None, bytebuffer))) when(ts.resourceOffers(any[IndexedSeq[WorkerOffer]], any[Boolean])).thenReturn(taskDescs) backend.driverEndpoint.send(ReviveOffers) @@ -606,6 +606,76 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo assert(mockEndpointRef.decommissionReceived) } + test("UpdateUserCredentials is broadcast to all registered executors") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + + // Register two mock executors + val mockEndpointRef1 = new MockExecutorRpcEndpointRef(conf) + val mockEndpointRef2 = new MockExecutorRpcEndpointRef(conf) + val mockAddress = mock[RpcAddress] + + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("1", mockEndpointRef1, mockAddress.host, 1, Map(), Map(), + Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("2", mockEndpointRef2, mockAddress.host, 1, Map(), Map(), + Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + + sc.listenerBus.waitUntilEmpty(executorUpTimeout.toMillis) + + // Neither executor should have received credentials yet + assert(mockEndpointRef1.receivedUserCredentials.isEmpty) + assert(mockEndpointRef2.receivedUserCredentials.isEmpty) + + // Send UpdateUserCredentials via DriverEndpoint + val testCredentials = Array[Byte](1, 2, 3, 4, 5) + backend.driverEndpoint.send(UpdateUserCredentials(testCredentials)) + + // Wait for the message to be processed + eventually(timeout(5 seconds)) { + assert(mockEndpointRef1.receivedUserCredentials.isDefined) + assert(mockEndpointRef2.receivedUserCredentials.isDefined) + } + + assert(mockEndpointRef1.receivedUserCredentials.get === testCredentials) + assert(mockEndpointRef2.receivedUserCredentials.get === testCredentials) + + // Verify SparkEnv credential store is also updated + assert(SparkEnv.get.userCredentials.get() === testCredentials) + } + + test("SparkAppConfig includes current user credentials for late-registering executors") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + + // Simulate credential acquisition by setting credentials before any executor registers + val testCredentials = Array[Byte](10, 20, 30, 40, 50) + backend.driverEndpoint.send(UpdateUserCredentials(testCredentials)) + + // Wait for the message to be processed + eventually(timeout(5 seconds)) { + assert(SparkEnv.get.userCredentials.get() != null) + } + + // Now retrieve SparkAppConfig as a late-registering executor would + val appConfig = backend.driverEndpoint.askSync[SparkAppConfig]( + RetrieveSparkAppConfig(ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + + // Verify that userCredentials is present in the response + assert(appConfig.userCredentials.isDefined, + "SparkAppConfig should include user credentials for late-registering executors") + assert(appConfig.userCredentials.get === testCredentials) + } + private def testSubmitJob(sc: SparkContext, rdd: RDD[Int]): Unit = { sc.submitJob( rdd, @@ -667,12 +737,15 @@ private[spark] class MockExecutorRpcEndpointRef(conf: SparkConf) extends RpcEndp // scalastyle:on executioncontextglobal var decommissionReceived = false + var receivedUserCredentials: Option[Array[Byte]] = None override def address: RpcAddress = null override def name: String = "executor" override def send(message: Any): Unit = message match { case DecommissionExecutor => decommissionReceived = true + case UpdateUserCredentials(creds) => receivedUserCredentials = Some(creds) + case _ => } override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = { Future{true.asInstanceOf[T]} diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala index dcb22f8b0d7d8..4880fa0e6bbbe 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala @@ -85,6 +85,7 @@ class TaskDescriptionSuite extends SparkFunSuite { originalProperties, cpus = 2, originalResources, + None, taskBuffer ) @@ -102,6 +103,7 @@ class TaskDescriptionSuite extends SparkFunSuite { assert(decodedTaskDescription.properties.equals(originalTaskDescription.properties)) assert(decodedTaskDescription.cpus.equals(originalTaskDescription.cpus)) assert(decodedTaskDescription.resources === originalTaskDescription.resources) + assert(decodedTaskDescription.userCredentials === None) assert(decodedTaskDescription.serializedTask.equals(taskBuffer)) } @@ -119,6 +121,7 @@ class TaskDescriptionSuite extends SparkFunSuite { new Properties(), cpus, Map.empty, + None, ByteBuffer.wrap(Array[Byte](1, 2, 3, 4))) val decoded = TaskDescription.decode(TaskDescription.encode(taskDescription)) // The decoded amount is what the executor echoes back in its status update and what the @@ -129,4 +132,65 @@ class TaskDescriptionSuite extends SparkFunSuite { } } + test("encoding and decoding TaskDescription with userCredentials preserves credentials") { + val taskBuffer = ByteBuffer.wrap(Array[Byte](1, 2, 3, 4)) + val properties = new Properties() + properties.put("key", "value") + + val credentialBytes = Array[Byte](10, 20, 30, 40, 50) + + val taskDescription = new TaskDescription( + taskId = 42, + attemptNumber = 0, + executorId = "exec-1", + name = "task with credentials", + index = 0, + partitionId = 0, + JobArtifactSet.emptyJobArtifactSet, + properties, + cpus = 1, + Map.empty, + Some(credentialBytes), + taskBuffer + ) + + val serialized = TaskDescription.encode(taskDescription) + val decoded = TaskDescription.decode(serialized) + + assert(decoded.taskId === 42) + assert(decoded.name === "task with credentials") + assert(decoded.userCredentials.isDefined) + assert(decoded.userCredentials.get === credentialBytes) + assert(decoded.serializedTask.equals(taskBuffer)) + } + + test("encoding and decoding TaskDescription without userCredentials round-trips correctly") { + val taskBuffer = ByteBuffer.wrap(Array[Byte](5, 6, 7, 8)) + val properties = new Properties() + + val taskDescription = new TaskDescription( + taskId = 99, + attemptNumber = 1, + executorId = "exec-2", + name = "task without credentials", + index = 3, + partitionId = 2, + JobArtifactSet.emptyJobArtifactSet, + properties, + cpus = 2, + Map.empty, + None, + taskBuffer + ) + + val serialized = TaskDescription.encode(taskDescription) + val decoded = TaskDescription.decode(serialized) + + assert(decoded.taskId === 99) + assert(decoded.attemptNumber === 1) + assert(decoded.name === "task without credentials") + assert(decoded.userCredentials === None) + assert(decoded.serializedTask.equals(taskBuffer)) + } + } diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala index e44d7e29ef606..787ab75fdec57 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala @@ -128,6 +128,11 @@ private[spark] object KubernetesExecutorBackend extends Logging { val env = SparkEnv.createExecutorEnv(driverConf, execId, arguments.bindAddress, arguments.hostname, arguments.cores, cfg.ioEncryptionKey, isLocal = false) + // Apply initial user credentials to the executor credential store. + cfg.userCredentials.foreach { credentials => + env.userCredentials.set(credentials) + } + val backend = backendCreateFn(env.rpcEnv, arguments, env, cfg.resourceProfile, execId) env.rpcEnv.setupEndpoint("Executor", backend) arguments.workerUrl.foreach { url =>