From e39bddab911f9eab824de95c67ff12bd86b87d00 Mon Sep 17 00:00:00 2001 From: fwang12 Date: Thu, 23 Mar 2023 13:00:49 +0800 Subject: [PATCH] [KYUUBI #4541] Support to customize the attributes to expose for Spark engine ### _Why are the changes needed?_ ### _How was this patch tested?_ - [ ] 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.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request Closes #4541 from turboFei/expose. Closes #4541 f882b4dda [fwang12] engine register attributes Authored-by: fwang12 Signed-off-by: fwang12 --- .../engine/spark/SparkTBinaryFrontendService.scala | 7 +++++-- .../main/scala/org/apache/kyuubi/config/KyuubiConf.scala | 9 +++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkTBinaryFrontendService.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkTBinaryFrontendService.scala index 49cb9b3efcd..854a28e85a1 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkTBinaryFrontendService.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkTBinaryFrontendService.scala @@ -27,6 +27,7 @@ import org.apache.spark.SparkContext import org.apache.spark.kyuubi.SparkContextHelper import org.apache.kyuubi.{KyuubiSQLException, Logging} +import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiReservedKeys._ import org.apache.kyuubi.ha.client.{EngineServiceDiscovery, ServiceDiscovery} import org.apache.kyuubi.service.{Serverable, Service, TBinaryFrontendService} @@ -94,8 +95,10 @@ class SparkTBinaryFrontendService( } override def attributes: Map[String, String] = { - val attributes = Map( - KYUUBI_ENGINE_ID -> KyuubiSparkUtil.engineId) + val extraAttributes = conf.get(KyuubiConf.ENGINE_SPARK_REGISTER_ATTRIBUTES).map { attr => + attr -> KyuubiSparkUtil.globalSparkContext.getConf.get(attr, "") + }.toMap + val attributes = extraAttributes ++ Map(KYUUBI_ENGINE_ID -> KyuubiSparkUtil.engineId) // TODO Support Spark Web UI Enabled SSL sc.uiWebUrl match { case Some(url) => attributes ++ Map(KYUUBI_ENGINE_URL -> url.split("//").last) diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala index 5a576a819a9..5e775595257 100644 --- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala +++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala @@ -2764,6 +2764,15 @@ object KyuubiConf { .stringConf .createWithDefault("bin/python") + val ENGINE_SPARK_REGISTER_ATTRIBUTES: ConfigEntry[Seq[String]] = + buildConf("kyuubi.engine.spark.register.attributes") + .internal + .doc("The extra attributes to expose when registering for Spark engine.") + .version("1.8.0") + .stringConf + .toSequence() + .createWithDefault(Seq("spark.driver.memory", "spark.executor.memory")) + val ENGINE_HIVE_EVENT_LOGGERS: ConfigEntry[Seq[String]] = buildConf("kyuubi.engine.hive.event.loggers") .doc("A comma-separated list of engine history loggers, where engine/session/operation etc" +