Skip to content

Commit 93bf260

Browse files
zwangshengpan3793
authored andcommitted
[KYUUBI #3527] [SPARK] Kyuubi should set env SPARK_USER_NAME for K8s deployment
### _Why are the changes needed?_ On case spark on kubernetes, spark using env `SPARK_USER_NAME` as user name. So kyuubi should build spark engine with this env when proxy user or using keytab. This conf only affect on kubernetes case. Ref: apache/spark#23017 ### _How was this patch tested?_ - [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible - [ ] Add screenshots for manual tests if appropriate - [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request Closes #3527 from zwangsheng/feature/add_spark_user_name. Closes #3527 9596372 [zwangsheng] only k8s case ddd713f [zwangsheng] fix 48b9b22 [zwangsheng] add Authored-by: zwangsheng <2213335496@qq.com> Signed-off-by: Cheng Pan <chengpan@apache.org> (cherry picked from commit 3738512) Signed-off-by: Cheng Pan <chengpan@apache.org>
1 parent 811f06b commit 93bf260

File tree

2 files changed

+55
-10
lines changed

2 files changed

+55
-10
lines changed

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,26 @@ class SparkProcessBuilder(
9191
buffer += s"${convertConfigKey(k)}=$v"
9292
}
9393

94-
// iff the keytab is specified, PROXY_USER is not supported
95-
if (!useKeytab()) {
96-
buffer += PROXY_USER
97-
buffer += proxyUser
94+
// For spark on kubernetes, spark pod using env SPARK_USER_NAME as current user
95+
def setSparkUserName(userName: String): Unit = {
96+
clusterManager().foreach(cm => {
97+
if (cm.startsWith("k8s://")) {
98+
buffer += CONF
99+
buffer += s"spark.kubernetes.driverEnv.SPARK_USER_NAME=$userName"
100+
buffer += CONF
101+
buffer += s"spark.kubernetes.executorEnv.SPARK_USER_NAME=$userName"
102+
}
103+
})
104+
}
105+
106+
// if the keytab is specified, PROXY_USER is not supported
107+
tryKeytab() match {
108+
case None =>
109+
setSparkUserName(proxyUser)
110+
buffer += PROXY_USER
111+
buffer += proxyUser
112+
case Some(name) =>
113+
setSparkUserName(name)
98114
}
99115

100116
mainResource.foreach { r => buffer += r }
@@ -104,26 +120,27 @@ class SparkProcessBuilder(
104120

105121
override protected def module: String = "kyuubi-spark-sql-engine"
106122

107-
private def useKeytab(): Boolean = {
123+
private def tryKeytab(): Option[String] = {
108124
val principal = conf.getOption(PRINCIPAL)
109125
val keytab = conf.getOption(KEYTAB)
110126
if (principal.isEmpty || keytab.isEmpty) {
111-
false
127+
None
112128
} else {
113129
try {
114130
val ugi = UserGroupInformation
115131
.loginUserFromKeytabAndReturnUGI(principal.get, keytab.get)
116-
val keytabEnabled = ugi.getShortUserName == proxyUser
117-
if (!keytabEnabled) {
132+
if (ugi.getShortUserName != proxyUser) {
118133
warn(s"The session proxy user: $proxyUser is not same with " +
119134
s"spark principal: ${ugi.getShortUserName}, so we can't support use keytab. " +
120135
s"Fallback to use proxy user.")
136+
None
137+
} else {
138+
Some(ugi.getShortUserName)
121139
}
122-
keytabEnabled
123140
} catch {
124141
case e: IOException =>
125142
error(s"Failed to login for ${principal.get}", e)
126-
false
143+
None
127144
}
128145
}
129146
}

kyuubi-server/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,34 @@ class SparkProcessBuilderSuite extends KerberizedTestHelper with MockitoSugar {
270270
assert(!pb.toString.contains(engineRefId2))
271271
assert(pb.toString.contains(engineRefId))
272272
}
273+
274+
test("SparkProcessBuilder build spark engine with SPARK_USER_NAME") {
275+
val proxyName = "kyuubi"
276+
val conf1 = KyuubiConf(false).set("spark.master", "k8s://test:12345")
277+
val b1 = new SparkProcessBuilder(proxyName, conf1)
278+
val c1 = b1.toString.split(' ')
279+
assert(c1.contains(s"spark.kubernetes.driverEnv.SPARK_USER_NAME=$proxyName"))
280+
assert(c1.contains(s"spark.kubernetes.executorEnv.SPARK_USER_NAME=$proxyName"))
281+
282+
tryWithSecurityEnabled {
283+
val conf2 = conf.set("spark.master", "k8s://test:12345")
284+
.set("spark.kerberos.principal", testPrincipal)
285+
.set("spark.kerberos.keytab", testKeytab)
286+
val name = ServiceUtils.getShortName(testPrincipal)
287+
val b2 = new SparkProcessBuilder(name, conf2)
288+
val c2 = b2.toString.split(' ')
289+
assert(c2.contains(s"spark.kubernetes.driverEnv.SPARK_USER_NAME=$name"))
290+
assert(c2.contains(s"spark.kubernetes.executorEnv.SPARK_USER_NAME=$name"))
291+
assert(!c2.contains(s"--proxy-user $name"))
292+
}
293+
294+
// Test no-kubernetes case
295+
val conf3 = KyuubiConf(false)
296+
val b3 = new SparkProcessBuilder(proxyName, conf3)
297+
val c3 = b3.toString.split(' ')
298+
assert(!c3.contains(s"spark.kubernetes.driverEnv.SPARK_USER_NAME=$proxyName"))
299+
assert(!c3.contains(s"spark.kubernetes.executorEnv.SPARK_USER_NAME=$proxyName"))
300+
}
273301
}
274302

275303
class FakeSparkProcessBuilder(config: KyuubiConf)

0 commit comments

Comments
 (0)